Skip to content

Commit

Permalink
downcast geometry array (#918)
Browse files Browse the repository at this point in the history 
todo:

- [x] Change downcast to be implemented in terms of `resolve_types` and
`cast`.
- [x] Handle slicing in `GeometryArray::has_points`. I.e. bring back the
slice offset and length. Then if the array has been sliced and the point
array exists, check if the point values are within the bounds of the
current slice. The fastest way to do this should be to check the length
of the point array, then as you're iterating through the type_ids array,
once you've seen that many ids pointing to the point array, you know the
array is fully there. Or, in the cast of downcasting you really only
care about whether _any_ geometry exists. So you can short-circuit as
soon as you've seen the first point.
- [x] Ensure we apply slicing when exporting the GeometryArray. E.g. in
the `TryFrom` impl from GeometryArray to `PointArray`, we need to call
the slice when we export.
- [ ] recursively expand the types from the geometry collection array.
But only if the geometry collection array has only a single geometry per
row and it could be split up.

Closes #416
  • Loading branch information
@kylebarron
kylebarron authored Dec 10, 2024
1 parent e4add1d commit 9e68c71
Show file tree
Hide file tree
Showing 13 changed files with 540 additions and 484 deletions.
107 changes: 44 additions & 63 deletions rust/geoarrow/src/algorithm/native/cast.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ impl Default for CastOptions {

/// Note: not currently used and outdated
#[allow(dead_code)]
fn can_cast_types(from_type: &NativeType, to_type: &NativeType) -> bool {
fn can_cast_types(from_type: NativeType, to_type: NativeType) -> bool {
if from_type == to_type {
return true;
}
Expand All @@ -51,13 +51,13 @@ pub trait Cast {
type Output;

/// Note: **does not currently implement dimension casts**
fn cast(&self, to_type: &NativeType) -> Self::Output;
fn cast(&self, to_type: NativeType) -> Self::Output;
}

impl Cast for PointArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -66,6 +66,7 @@ impl Cast for PointArray {
MultiPoint(_, _) => Ok(Arc::new(MultiPointArray::from(array))),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -76,7 +77,7 @@ impl Cast for PointArray {
impl Cast for LineStringArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -86,6 +87,7 @@ impl Cast for LineStringArray {
MultiLineString(_, _) => Ok(Arc::new(MultiLineStringArray::from(array))),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -96,7 +98,7 @@ impl Cast for LineStringArray {
impl Cast for PolygonArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -106,6 +108,7 @@ impl Cast for PolygonArray {
MultiPolygon(_, _) => Ok(Arc::new(MultiPolygonArray::from(array))),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -116,7 +119,7 @@ impl Cast for PolygonArray {
impl Cast for MultiPointArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -126,6 +129,7 @@ impl Cast for MultiPointArray {
MultiPoint(_, _) => Ok(Arc::new(array)),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -136,7 +140,7 @@ impl Cast for MultiPointArray {
impl Cast for MultiLineStringArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -145,6 +149,7 @@ impl Cast for MultiLineStringArray {
LineString(_, _) => Ok(Arc::new(LineStringArray::try_from(array)?)),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -155,7 +160,7 @@ impl Cast for MultiLineStringArray {
impl Cast for MultiPolygonArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -164,6 +169,7 @@ impl Cast for MultiPolygonArray {
Polygon(_, _) => Ok(Arc::new(PolygonArray::try_from(array)?)),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::from(array))),
GeometryCollection(_, _) => Ok(Arc::new(GeometryCollectionArray::from(array))),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
Expand All @@ -174,7 +180,7 @@ impl Cast for MultiPolygonArray {
impl Cast for MixedGeometryArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -198,7 +204,7 @@ impl Cast for MixedGeometryArray {
impl Cast for GeometryCollectionArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
use NativeType::*;

let array = self.to_coord_type(to_type.coord_type());
Expand All @@ -212,17 +218,29 @@ impl Cast for GeometryCollectionArray {
MultiPolygon(_, _) => Ok(Arc::new(MultiPolygonArray::try_from(array)?)),
Mixed(_, _) => Ok(Arc::new(MixedGeometryArray::try_from(array)?)),
GeometryCollection(_, _) => Ok(Arc::new(array)),
Geometry(_) => Ok(Arc::new(GeometryArray::from(array))),
dt => Err(GeoArrowError::General(format!(
"invalid cast to type {dt:?}"
))),
}
}
}

impl Cast for GeometryArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: NativeType) -> Self::Output {
// TODO: validate dimension
let array = self.to_coord_type(to_type.coord_type());
let mixed_array = MixedGeometryArray::try_from(array)?;
mixed_array.cast(to_type)
}
}

impl Cast for &dyn NativeArray {
type Output = Result<Arc<dyn NativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
// TODO: not working :/
// if self.data_type() == to_type {
// return Ok(Arc::new(self.to_owned()));
Expand All @@ -239,56 +257,18 @@ impl Cast for &dyn NativeArray {
MultiPolygon(_, _) => self.as_ref().as_multi_polygon().cast(to_type),
Mixed(_, _) => self.as_ref().as_mixed().cast(to_type),
GeometryCollection(_, _) => self.as_ref().as_geometry_collection().cast(to_type),
Geometry(_) => self.as_ref().as_geometry().cast(to_type),
_ => todo!(),
}
}
}

macro_rules! impl_chunked_cast_non_generic {
($chunked_array:ty) => {
impl Cast for $chunked_array {
type Output = Result<Arc<dyn ChunkedNativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
macro_rules! impl_cast {
($method:ident) => {
Arc::new(ChunkedGeometryArray::new(
self.geometry_chunks()
.iter()
.map(|chunk| {
Ok(chunk.as_ref().cast(to_type)?.as_ref().$method().clone())
})
.collect::<Result<Vec<_>>>()?,
))
};
}

use NativeType::*;

let result: Arc<dyn ChunkedNativeArray> = match to_type {
Point(_, _) => impl_cast!(as_point),
LineString(_, _) => impl_cast!(as_line_string),
Polygon(_, _) => impl_cast!(as_polygon),
MultiPoint(_, _) => impl_cast!(as_multi_point),
MultiLineString(_, _) => impl_cast!(as_multi_line_string),
MultiPolygon(_, _) => impl_cast!(as_multi_polygon),
Mixed(_, _) => impl_cast!(as_mixed),
GeometryCollection(_, _) => impl_cast!(as_geometry_collection),
Rect(_) => impl_cast!(as_rect),
Geometry(_) => todo!("cast to unknown"),
};
Ok(result)
}
}
};
}

macro_rules! impl_chunked_cast_generic {
macro_rules! impl_chunked_cast {
($chunked_array:ty) => {
impl Cast for $chunked_array {
type Output = Result<Arc<dyn ChunkedNativeArray>>;

fn cast(&self, to_type: &NativeType) -> Self::Output {
fn cast(&self, to_type: NativeType) -> Self::Output {
macro_rules! impl_cast {
($method:ident) => {
Arc::new(ChunkedGeometryArray::new(
Expand All @@ -314,21 +294,22 @@ macro_rules! impl_chunked_cast_generic {
Mixed(_, _) => impl_cast!(as_mixed),
GeometryCollection(_, _) => impl_cast!(as_geometry_collection),
Rect(_) => impl_cast!(as_rect),
Geometry(_) => todo!("cast to unknown"),
Geometry(_) => impl_cast!(as_geometry),
};
Ok(result)
}
}
};
}

impl_chunked_cast_non_generic!(ChunkedPointArray);
impl_chunked_cast_non_generic!(ChunkedRectArray);
impl_chunked_cast_non_generic!(&dyn ChunkedNativeArray);
impl_chunked_cast_generic!(ChunkedLineStringArray);
impl_chunked_cast_generic!(ChunkedPolygonArray);
impl_chunked_cast_generic!(ChunkedMultiPointArray);
impl_chunked_cast_generic!(ChunkedMultiLineStringArray);
impl_chunked_cast_generic!(ChunkedMultiPolygonArray);
impl_chunked_cast_generic!(ChunkedMixedGeometryArray);
impl_chunked_cast_generic!(ChunkedGeometryCollectionArray);
impl_chunked_cast!(ChunkedPointArray);
impl_chunked_cast!(ChunkedRectArray);
impl_chunked_cast!(&dyn ChunkedNativeArray);
impl_chunked_cast!(ChunkedLineStringArray);
impl_chunked_cast!(ChunkedPolygonArray);
impl_chunked_cast!(ChunkedMultiPointArray);
impl_chunked_cast!(ChunkedMultiLineStringArray);
impl_chunked_cast!(ChunkedMultiPolygonArray);
impl_chunked_cast!(ChunkedMixedGeometryArray);
impl_chunked_cast!(ChunkedGeometryCollectionArray);
impl_chunked_cast!(ChunkedUnknownGeometryArray);
Loading

0 comments on commit 9e68c71

Please sign in to comment.