Skip to content

Commit

Permalink
Fix GeometryArray FFI
Browse files Browse the repository at this point in the history
  • Loading branch information
@kylebarron
kylebarron committed Dec 20, 2024
1 parent 19d8715 commit 910d92b
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 52 deletions.
7 changes: 3 additions & 4 deletions python/tests/interop/test_wkb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,9 @@ def test_geometry_collection():
retour = to_wkb(parsed_geoarrow)
retour_shapely = shapely.from_wkb(retour[0].as_py())

# Need to unpack the geoms because they're returned as multi-geoms
assert retour_shapely.geoms[0].geoms[0] == point
assert retour_shapely.geoms[1].geoms[0] == point2
assert retour_shapely.geoms[2].geoms[0] == line_string
assert retour_shapely.geoms[0] == point
assert retour_shapely.geoms[1] == point2
assert retour_shapely.geoms[2] == line_string


def test_ewkb_srid():
Expand Down
18 changes: 6 additions & 12 deletions rust/geoarrow/src/array/geometry/array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -676,18 +676,14 @@ impl NativeGeometryAccessor for GeometryArray {
4 => Geometry::MultiPoint(self.mpoint_xy.value(offset)),
5 => Geometry::MultiLineString(self.mline_string_xy.value(offset)),
6 => Geometry::MultiPolygon(self.mpolygon_xy.value(offset)),
7 => {
panic!("nested geometry collections not supported")
}
7 => Geometry::GeometryCollection(self.gc_xy.value(offset)),
11 => Geometry::Point(self.point_xyz.value(offset)),
12 => Geometry::LineString(self.line_string_xyz.value(offset)),
13 => Geometry::Polygon(self.polygon_xyz.value(offset)),
14 => Geometry::MultiPoint(self.mpoint_xyz.value(offset)),
15 => Geometry::MultiLineString(self.mline_string_xyz.value(offset)),
16 => Geometry::MultiPolygon(self.mpolygon_xyz.value(offset)),
17 => {
panic!("nested geometry collections not supported")
}
17 => Geometry::GeometryCollection(self.gc_xyz.value(offset)),
_ => panic!("unknown type_id {}", type_id),
}
}
Expand Down Expand Up @@ -719,18 +715,14 @@ impl<'a> ArrayAccessor<'a> for GeometryArray {
4 => Geometry::MultiPoint(self.mpoint_xy.value(offset)),
5 => Geometry::MultiLineString(self.mline_string_xy.value(offset)),
6 => Geometry::MultiPolygon(self.mpolygon_xy.value(offset)),
7 => {
panic!("nested geometry collections not supported")
}
7 => Geometry::GeometryCollection(self.gc_xy.value(offset)),
11 => Geometry::Point(self.point_xyz.value(offset)),
12 => Geometry::LineString(self.line_string_xyz.value(offset)),
13 => Geometry::Polygon(self.polygon_xyz.value(offset)),
14 => Geometry::MultiPoint(self.mpoint_xyz.value(offset)),
15 => Geometry::MultiLineString(self.mline_string_xyz.value(offset)),
16 => Geometry::MultiPolygon(self.mpolygon_xyz.value(offset)),
17 => {
panic!("nested geometry collections not supported")
}
17 => Geometry::GeometryCollection(self.gc_xyz.value(offset)),
_ => panic!("unknown type_id {}", type_id),
}
}
Expand All @@ -752,12 +744,14 @@ impl IntoArrow for GeometryArray {
self.mpoint_xy.into_array_ref(),
self.mline_string_xy.into_array_ref(),
self.mpolygon_xy.into_array_ref(),
self.gc_xy.into_array_ref(),
self.point_xyz.into_array_ref(),
self.line_string_xyz.into_array_ref(),
self.polygon_xyz.into_array_ref(),
self.mpoint_xyz.into_array_ref(),
self.mline_string_xyz.into_array_ref(),
self.mpolygon_xyz.into_array_ref(),
self.gc_xyz.into_array_ref(),
];

UnionArray::try_new(
Expand Down
96 changes: 63 additions & 33 deletions rust/geoarrow/src/datatypes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,9 +333,9 @@ fn rect_data_type(dim: Dimension) -> DataType {
DataType::Struct(rect_fields(dim))
}

fn unknown_data_type(coord_type: CoordType) -> DataType {
fn geometry_data_type(coord_type: CoordType) -> DataType {
let mut fields = vec![];
let type_ids = vec![1, 2, 3, 4, 5, 6, 11, 12, 13, 14, 15, 16];
let type_ids = vec![1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17];

// Note: we manually construct the fields because these fields shouldn't have their own
// GeoArrow extension metadata
Expand All @@ -344,42 +344,72 @@ fn unknown_data_type(coord_type: CoordType) -> DataType {
NativeType::Point(coord_type, Dimension::XY).to_data_type(),
true,
));

let linestring = NativeType::LineString(coord_type, Dimension::XY);
fields.push(Field::new("", linestring.to_data_type(), true));

let polygon = NativeType::Polygon(coord_type, Dimension::XY);
fields.push(Field::new("", polygon.to_data_type(), true));

let multi_point = NativeType::MultiPoint(coord_type, Dimension::XY);
fields.push(Field::new("", multi_point.to_data_type(), true));

let multi_line_string = NativeType::MultiLineString(coord_type, Dimension::XY);
fields.push(Field::new("", multi_line_string.to_data_type(), true));

let multi_polygon = NativeType::MultiPolygon(coord_type, Dimension::XY);
fields.push(Field::new("", multi_polygon.to_data_type(), true));
fields.push(Field::new(
"",
NativeType::LineString(coord_type, Dimension::XY).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::Polygon(coord_type, Dimension::XY).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiPoint(coord_type, Dimension::XY).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiLineString(coord_type, Dimension::XY).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiPolygon(coord_type, Dimension::XY).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::GeometryCollection(coord_type, Dimension::XY).to_data_type(),
true,
));

fields.push(Field::new(
"",
NativeType::Point(coord_type, Dimension::XYZ).to_data_type(),
true,
));

let linestring = NativeType::LineString(coord_type, Dimension::XYZ);
fields.push(Field::new("", linestring.to_data_type(), true));

let polygon = NativeType::Polygon(coord_type, Dimension::XYZ);
fields.push(Field::new("", polygon.to_data_type(), true));

let multi_point = NativeType::MultiPoint(coord_type, Dimension::XYZ);
fields.push(Field::new("", multi_point.to_data_type(), true));

let multi_line_string = NativeType::MultiLineString(coord_type, Dimension::XYZ);
fields.push(Field::new("", multi_line_string.to_data_type(), true));

let multi_polygon = NativeType::MultiPolygon(coord_type, Dimension::XYZ);
fields.push(Field::new("", multi_polygon.to_data_type(), true));
fields.push(Field::new(
"",
NativeType::LineString(coord_type, Dimension::XYZ).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::Polygon(coord_type, Dimension::XYZ).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiPoint(coord_type, Dimension::XYZ).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiLineString(coord_type, Dimension::XYZ).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::MultiPolygon(coord_type, Dimension::XYZ).to_data_type(),
true,
));
fields.push(Field::new(
"",
NativeType::GeometryCollection(coord_type, Dimension::XYZ).to_data_type(),
true,
));

let union_fields = UnionFields::new(type_ids, fields);
DataType::Union(union_fields, UnionMode::Dense)
Expand Down Expand Up @@ -445,7 +475,7 @@ impl NativeType {
MultiPolygon(coord_type, dim) => multi_polygon_data_type(*coord_type, *dim),
GeometryCollection(coord_type, dim) => geometry_collection_data_type(*coord_type, *dim),
Rect(dim) => rect_data_type(*dim),
Geometry(coord_type) => unknown_data_type(*coord_type),
Geometry(coord_type) => geometry_data_type(*coord_type),
}
}

Expand Down
6 changes: 3 additions & 3 deletions rust/geoarrow/src/io/wkb/api.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ impl FromWKB for MixedGeometryArray {
) -> Result<Self> {
let wkb_objects: Vec<Option<WKB<'_, O>>> = arr.iter().collect();
let builder =
MixedGeometryBuilder::from_wkb(&wkb_objects, dim, coord_type, arr.metadata(), true)?;
MixedGeometryBuilder::from_wkb(&wkb_objects, dim, coord_type, arr.metadata(), false)?;
Ok(builder.finish())
}
}
Expand All @@ -96,7 +96,7 @@ impl FromWKB for GeometryCollectionArray {
dim,
coord_type,
arr.metadata(),
true,
false,
)?;
Ok(builder.finish())
}
Expand All @@ -111,7 +111,7 @@ impl FromWKB for GeometryArray {
_dim: Dimension,
) -> Result<Self> {
let wkb_objects: Vec<Option<WKB<'_, O>>> = arr.iter().collect();
let builder = GeometryBuilder::from_wkb(&wkb_objects, coord_type, arr.metadata(), true)?;
let builder = GeometryBuilder::from_wkb(&wkb_objects, coord_type, arr.metadata(), false)?;
Ok(builder.finish())
}
}
Expand Down

0 comments on commit 910d92b

Please sign in to comment.