Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add geometry collection as children of unknown geometry array #896

Merged
merged 2 commits into from
Dec 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
184 changes: 85 additions & 99 deletions rust/geoarrow/src/algorithm/native/downcast.rs

Large diffs are not rendered by default.

26 changes: 15 additions & 11 deletions rust/geoarrow/src/array/geometrycollection/array.rs
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
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
Loading