Skip to content

Commit

Permalink
Add geometry collection as children of unknown geometry array
Browse files Browse the repository at this point in the history
  • Loading branch information
@kylebarron
kylebarron committed Dec 4, 2024
1 parent ba0bd65 commit 73aa2e4
Show file tree
Hide file tree
Showing 6 changed files with 368 additions and 138 deletions.
44 changes: 20 additions & 24 deletions rust/geoarrow/src/algorithm/native/downcast.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use crate::schema::GeoSchemaExt;
use crate::table::Table;
use crate::NativeArray;

/// Downcast will change between geometry types but will not affect the dimension of the data.
pub trait Downcast {
type Output;

Expand Down Expand Up @@ -100,10 +101,7 @@ impl Downcast for LineStringArray {

fn downcast(&self) -> Self::Output {
match (self.data_type(), self.downcasted_data_type()) {
(
NativeType::LineString(_, Dimension::XY),
NativeType::LineString(_, Dimension::XY),
) => Arc::new(self.clone()),
(NativeType::LineString(_, _), NativeType::LineString(_, _)) => Arc::new(self.clone()),
_ => unreachable!(),
}
}
Expand Down Expand Up @@ -381,37 +379,35 @@ impl Downcast for &dyn NativeArray {
type Output = Arc<dyn NativeArray>;

fn downcasted_data_type(&self) -> NativeType {
use Dimension::*;
use NativeType::*;

match self.data_type() {
Point(_, XY) => self.as_point().downcasted_data_type(),
LineString(_, XY) => self.as_line_string().downcasted_data_type(),
Polygon(_, XY) => self.as_polygon().downcasted_data_type(),
MultiPoint(_, XY) => self.as_multi_point().downcasted_data_type(),
MultiLineString(_, XY) => self.as_multi_line_string().downcasted_data_type(),
MultiPolygon(_, XY) => self.as_multi_polygon().downcasted_data_type(),
Mixed(_, XY) => self.as_mixed().downcasted_data_type(),
GeometryCollection(_, XY) => self.as_geometry_collection().downcasted_data_type(),
Rect(XY) => self.as_rect().downcasted_data_type(),
Point(_, _) => self.as_point().downcasted_data_type(),
LineString(_, _) => self.as_line_string().downcasted_data_type(),
Polygon(_, _) => self.as_polygon().downcasted_data_type(),
MultiPoint(_, _) => self.as_multi_point().downcasted_data_type(),
MultiLineString(_, _) => self.as_multi_line_string().downcasted_data_type(),
MultiPolygon(_, _) => self.as_multi_polygon().downcasted_data_type(),
Mixed(_, _) => self.as_mixed().downcasted_data_type(),
GeometryCollection(_, _) => self.as_geometry_collection().downcasted_data_type(),
Rect(_) => self.as_rect().downcasted_data_type(),
_ => todo!("3d support"),
}
}

fn downcast(&self) -> Self::Output {
use Dimension::*;
use NativeType::*;

match self.data_type() {
Point(_, XY) => self.as_point().downcast(),
LineString(_, XY) => self.as_line_string().downcast(),
Polygon(_, XY) => self.as_polygon().downcast(),
MultiPoint(_, XY) => self.as_multi_point().downcast(),
MultiLineString(_, XY) => self.as_multi_line_string().downcast(),
MultiPolygon(_, XY) => self.as_multi_polygon().downcast(),
Mixed(_, XY) => self.as_mixed().downcast(),
GeometryCollection(_, XY) => self.as_geometry_collection().downcast(),
Rect(XY) => self.as_rect().downcast(),
Point(_, _) => self.as_point().downcast(),
LineString(_, _) => self.as_line_string().downcast(),
Polygon(_, _) => self.as_polygon().downcast(),
MultiPoint(_, _) => self.as_multi_point().downcast(),
MultiLineString(_, _) => self.as_multi_line_string().downcast(),
MultiPolygon(_, _) => self.as_multi_polygon().downcast(),
Mixed(_, _) => self.as_mixed().downcast(),
GeometryCollection(_, _) => self.as_geometry_collection().downcast(),
Rect(_) => self.as_rect().downcast(),
_ => todo!("3d support"),
}
}
Expand Down
26 changes: 15 additions & 11 deletions rust/geoarrow/src/array/geometrycollection/array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -253,11 +253,11 @@ impl IntoArrow for GeometryCollectionArray {
}
}

impl TryFrom<&GenericListArray<i32>> for GeometryCollectionArray {
impl TryFrom<(&GenericListArray<i32>, Dimension)> for GeometryCollectionArray {
type Error = GeoArrowError;

fn try_from(value: &GenericListArray<i32>) -> Result<Self> {
let geoms: MixedGeometryArray = value.values().as_ref().try_into()?;
fn try_from((value, dim): (&GenericListArray<i32>, Dimension)) -> Result<Self> {
let geoms: MixedGeometryArray = (value.values().as_ref(), dim).try_into()?;
let geom_offsets = value.offsets();
let validity = value.nulls();

Expand All @@ -270,11 +270,11 @@ impl TryFrom<&GenericListArray<i32>> for GeometryCollectionArray {
}
}

impl TryFrom<&GenericListArray<i64>> for GeometryCollectionArray {
impl TryFrom<(&GenericListArray<i64>, Dimension)> for GeometryCollectionArray {
type Error = GeoArrowError;

fn try_from(value: &GenericListArray<i64>) -> Result<Self> {
let geoms: MixedGeometryArray = value.values().as_ref().try_into()?;
fn try_from((value, dim): (&GenericListArray<i64>, Dimension)) -> Result<Self> {
let geoms: MixedGeometryArray = (value.values().as_ref(), dim).try_into()?;
let geom_offsets = offsets_buffer_i64_to_i32(value.offsets())?;
let validity = value.nulls();

Expand All @@ -287,18 +287,18 @@ impl TryFrom<&GenericListArray<i64>> for GeometryCollectionArray {
}
}

impl TryFrom<&dyn Array> for GeometryCollectionArray {
impl TryFrom<(&dyn Array, Dimension)> for GeometryCollectionArray {
type Error = GeoArrowError;

fn try_from(value: &dyn Array) -> Result<Self> {
fn try_from((value, dim): (&dyn Array, Dimension)) -> Result<Self> {
match value.data_type() {
DataType::List(_) => {
let downcasted = value.as_list::<i32>();
downcasted.try_into()
(downcasted, dim).try_into()
}
DataType::LargeList(_) => {
let downcasted = value.as_list::<i64>();
downcasted.try_into()
(downcasted, dim).try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
Expand All @@ -312,7 +312,11 @@ impl TryFrom<(&dyn Array, &Field)> for GeometryCollectionArray {
type Error = GeoArrowError;

fn try_from((arr, field): (&dyn Array, &Field)) -> Result<Self> {
let mut arr: Self = arr.try_into()?;
let geom_type = NativeType::try_from(field)?;
let dim = geom_type
.dimension()
.ok_or(GeoArrowError::General("Expected dimension".to_string()))?;
let mut arr: Self = (arr, dim).try_into()?;
arr.metadata = Arc::new(ArrayMetadata::try_from(field)?);
Ok(arr)
}
Expand Down
104 changes: 43 additions & 61 deletions rust/geoarrow/src/array/mixed/array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -484,24 +484,25 @@ impl IntoArrow for MixedGeometryArray {
}
}

impl TryFrom<&UnionArray> for MixedGeometryArray {
impl TryFrom<(&UnionArray, Dimension)> for MixedGeometryArray {
type Error = GeoArrowError;

fn try_from(value: &UnionArray) -> std::result::Result<Self, Self::Error> {
let mut points: Vec<PointArray> = vec![];
let mut line_strings: Vec<LineStringArray> = vec![];
let mut polygons: Vec<PolygonArray> = vec![];
let mut multi_points: Vec<MultiPointArray> = vec![];
let mut multi_line_strings: Vec<MultiLineStringArray> = vec![];
let mut multi_polygons: Vec<MultiPolygonArray> = vec![];
fn try_from((value, dim): (&UnionArray, Dimension)) -> std::result::Result<Self, Self::Error> {
let mut points: Option<PointArray> = None;
let mut line_strings: Option<LineStringArray> = None;
let mut polygons: Option<PolygonArray> = None;
let mut multi_points: Option<MultiPointArray> = None;
let mut multi_line_strings: Option<MultiLineStringArray> = None;
let mut multi_polygons: Option<MultiPolygonArray> = None;

match value.data_type() {
DataType::Union(fields, mode) => {
if !matches!(mode, UnionMode::Dense) {
return Err(GeoArrowError::General("Expected dense union".to_string()));
}

for (type_id, _field) in fields.iter() {
let dimension = if type_id < 10 {
let found_dimension = if type_id < 10 {
Dimension::XY
} else if type_id < 20 {
Dimension::XYZ
Expand All @@ -512,48 +513,33 @@ impl TryFrom<&UnionArray> for MixedGeometryArray {
)));
};

if dim != found_dimension {
return Err( GeoArrowError::General(format!("expected dimension: {:?}, found child array with dimension {:?} and type_id: {}", dim, found_dimension, type_id )));
}

match type_id {
1 | 11 => {
points.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
points = Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
2 | 12 => {
line_strings.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
line_strings =
Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
3 | 13 => {
polygons.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
polygons =
Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
4 | 14 => {
multi_points.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
multi_points =
Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
5 | 15 => {
multi_line_strings.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
multi_line_strings =
Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
6 | 16 => {
multi_polygons.push(
(value.child(type_id).as_ref(), dimension)
.try_into()
.unwrap(),
);
multi_polygons =
Some((value.child(type_id).as_ref(), dim).try_into().unwrap());
}
_ => {
return Err(GeoArrowError::General(format!(
Expand All @@ -571,37 +557,28 @@ impl TryFrom<&UnionArray> for MixedGeometryArray {
// This is after checking for dense union
let offsets = value.offsets().unwrap().clone();

// TODO: make nicer errors. We don't currently allow a mixed geometry array with multiple
// dimensions of underlying geometries.
assert!(points.len() <= 1);
assert!(line_strings.len() <= 1);
assert!(polygons.len() <= 1);
assert!(multi_points.len() <= 1);
assert!(multi_line_strings.len() <= 1);
assert!(multi_polygons.len() <= 1);

Ok(Self::new(
type_ids,
offsets,
points.first().cloned().unwrap_or_default(),
line_strings.first().cloned().unwrap_or_default(),
polygons.first().cloned().unwrap_or_default(),
multi_points.first().cloned().unwrap_or_default(),
multi_line_strings.first().cloned().unwrap_or_default(),
multi_polygons.first().cloned().unwrap_or_default(),
points.unwrap_or_default(),
line_strings.unwrap_or_default(),
polygons.unwrap_or_default(),
multi_points.unwrap_or_default(),
multi_line_strings.unwrap_or_default(),
multi_polygons.unwrap_or_default(),
Default::default(),
))
}
}

impl TryFrom<&dyn Array> for MixedGeometryArray {
impl TryFrom<(&dyn Array, Dimension)> for MixedGeometryArray {
type Error = GeoArrowError;

fn try_from(value: &dyn Array) -> Result<Self> {
fn try_from((value, dim): (&dyn Array, Dimension)) -> Result<Self> {
match value.data_type() {
DataType::Union(_, _) => {
let downcasted = value.as_any().downcast_ref::<UnionArray>().unwrap();
downcasted.try_into()
(downcasted, dim).try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
Expand All @@ -611,11 +588,16 @@ impl TryFrom<&dyn Array> for MixedGeometryArray {
}
}

// TODO:, thinking all geoarrow.geometry will go through primary dimensionless geometry array
impl TryFrom<(&dyn Array, &Field)> for MixedGeometryArray {
type Error = GeoArrowError;

fn try_from((arr, field): (&dyn Array, &Field)) -> Result<Self> {
let mut arr: Self = arr.try_into()?;
let geom_type = NativeType::try_from(field)?;
let dim = geom_type
.dimension()
.ok_or(GeoArrowError::General("Expected dimension".to_string()))?;
let mut arr: Self = (arr, dim).try_into()?;
arr.metadata = Arc::new(ArrayMetadata::try_from(field)?);
Ok(arr)
}
Expand Down Expand Up @@ -871,7 +853,7 @@ mod test {

// Round trip to/from arrow-rs
let arrow_array = arr.into_arrow();
let round_trip_arr: MixedGeometryArray = (&arrow_array).try_into().unwrap();
let round_trip_arr: MixedGeometryArray = (&arrow_array, Dimension::XY).try_into().unwrap();

assert_eq!(
round_trip_arr.value_as_geo(0),
Expand Down Expand Up @@ -901,7 +883,7 @@ mod test {

// Round trip to/from arrow-rs
let arrow_array = arr.into_arrow();
let round_trip_arr: MixedGeometryArray = (&arrow_array).try_into().unwrap();
let round_trip_arr: MixedGeometryArray = (&arrow_array, Dimension::XY).try_into().unwrap();

assert_eq!(round_trip_arr.value_as_geo(0), geoms[0]);
assert_eq!(round_trip_arr.value_as_geo(1), geoms[1]);
Expand All @@ -918,7 +900,7 @@ mod test {

// Round trip to/from arrow-rs
let arrow_array = arr.into_arrow();
let round_trip_arr: MixedGeometryArray = (&arrow_array).try_into().unwrap();
let round_trip_arr: MixedGeometryArray = (&arrow_array, Dimension::XY).try_into().unwrap();

assert_eq!(round_trip_arr.value_as_geo(0), geoms[0]);
assert_eq!(round_trip_arr.value_as_geo(1), geoms[1]);
Expand Down
Loading

0 comments on commit 73aa2e4

Please sign in to comment.