Skip to content

Commit 0c1d120

Browse files
kylebarronkylebarron
authored
feat(geoarrow-schema): GeometryTypeId trait to infer GeoArrow type id for a type (#1372)
1 parent 274f27c commit 0c1d120

21 files changed

+310
-122
lines changed

‎rust/geoarrow-array/src/array/geometry.rs‎

Lines changed: 28 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use arrow_array::{Array, ArrayRef, OffsetSizeTrait, UnionArray};
66
use arrow_buffer::{NullBuffer, ScalarBuffer};
77
use arrow_schema::{ArrowError, DataType, Field, UnionMode};
88
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
9+
use geoarrow_schema::type_id::GeometryTypeId;
910
use geoarrow_schema::{
1011
CoordType, Dimension, GeoArrowType, GeometryCollectionType, GeometryType, LineStringType,
1112
Metadata, MultiLineStringType, MultiPointType, MultiPolygonType, PointType, PolygonType,
@@ -351,13 +352,13 @@ impl GeoArrowArray for GeometryArray {
351352
let offset = self.offsets[i] as usize;
352353
let dim = (type_id / 10) as usize;
353354
match type_id % 10 {
354-
1 => self.points[dim].is_null(offset),
355-
2 => self.line_strings[dim].is_null(offset),
356-
3 => self.polygons[dim].is_null(offset),
357-
4 => self.mpoints[dim].is_null(offset),
358-
5 => self.mline_strings[dim].is_null(offset),
359-
6 => self.mpolygons[dim].is_null(offset),
360-
7 => self.gcs[dim].is_null(offset),
355+
PointType::GEOMETRY_TYPE_OFFSET => self.points[dim].is_null(offset),
356+
LineStringType::GEOMETRY_TYPE_OFFSET => self.line_strings[dim].is_null(offset),
357+
PolygonType::GEOMETRY_TYPE_OFFSET => self.polygons[dim].is_null(offset),
358+
MultiPointType::GEOMETRY_TYPE_OFFSET => self.mpoints[dim].is_null(offset),
359+
MultiLineStringType::GEOMETRY_TYPE_OFFSET => self.mline_strings[dim].is_null(offset),
360+
MultiPolygonType::GEOMETRY_TYPE_OFFSET => self.mpolygons[dim].is_null(offset),
361+
GeometryCollectionType::GEOMETRY_TYPE_OFFSET => self.gcs[dim].is_null(offset),
361362
_ => unreachable!("unknown type_id {}", type_id),
362363
}
363364
}
@@ -385,13 +386,25 @@ impl<'a> GeoArrowArrayAccessor<'a> for GeometryArray {
385386
let dim = (type_id / 10) as usize;
386387

387388
let result = match type_id % 10 {
388-
1 => Geometry::Point(self.points[dim].value(offset)?),
389-
2 => Geometry::LineString(self.line_strings[dim].value(offset)?),
390-
3 => Geometry::Polygon(self.polygons[dim].value(offset)?),
391-
4 => Geometry::MultiPoint(self.mpoints[dim].value(offset)?),
392-
5 => Geometry::MultiLineString(self.mline_strings[dim].value(offset)?),
393-
6 => Geometry::MultiPolygon(self.mpolygons[dim].value(offset)?),
394-
7 => Geometry::GeometryCollection(self.gcs[dim].value(offset)?),
389+
PointType::GEOMETRY_TYPE_OFFSET => Geometry::Point(self.points[dim].value(offset)?),
390+
LineStringType::GEOMETRY_TYPE_OFFSET => {
391+
Geometry::LineString(self.line_strings[dim].value(offset)?)
392+
}
393+
PolygonType::GEOMETRY_TYPE_OFFSET => {
394+
Geometry::Polygon(self.polygons[dim].value(offset)?)
395+
}
396+
MultiPointType::GEOMETRY_TYPE_OFFSET => {
397+
Geometry::MultiPoint(self.mpoints[dim].value(offset)?)
398+
}
399+
MultiLineStringType::GEOMETRY_TYPE_OFFSET => {
400+
Geometry::MultiLineString(self.mline_strings[dim].value(offset)?)
401+
}
402+
MultiPolygonType::GEOMETRY_TYPE_OFFSET => {
403+
Geometry::MultiPolygon(self.mpolygons[dim].value(offset)?)
404+
}
405+
GeometryCollectionType::GEOMETRY_TYPE_OFFSET => {
406+
Geometry::GeometryCollection(self.gcs[dim].value(offset)?)
407+
}
395408
_ => unreachable!("unknown type_id {}", type_id),
396409
};
397410
Ok(result)
@@ -750,28 +763,6 @@ impl PartialEq for GeometryArray {
750763
}
751764
}
752765

753-
impl TypeId for PointArray {
754-
const ARRAY_TYPE_OFFSET: i8 = 1;
755-
}
756-
impl TypeId for LineStringArray {
757-
const ARRAY_TYPE_OFFSET: i8 = 2;
758-
}
759-
impl TypeId for PolygonArray {
760-
const ARRAY_TYPE_OFFSET: i8 = 3;
761-
}
762-
impl TypeId for MultiPointArray {
763-
const ARRAY_TYPE_OFFSET: i8 = 4;
764-
}
765-
impl TypeId for MultiLineStringArray {
766-
const ARRAY_TYPE_OFFSET: i8 = 5;
767-
}
768-
impl TypeId for MultiPolygonArray {
769-
const ARRAY_TYPE_OFFSET: i8 = 6;
770-
}
771-
impl TypeId for GeometryCollectionArray {
772-
const ARRAY_TYPE_OFFSET: i8 = 7;
773-
}
774-
775766
type ChildrenArrays = (
776767
[PointArray; 4],
777768
[LineStringArray; 4],
@@ -848,7 +839,7 @@ macro_rules! impl_primitive_cast {
848839
let dim = value.data_type.dimension();
849840
let metadata = value.data_type.metadata().clone();
850841

851-
let type_ids = vec![value.type_id(dim); value.len()].into();
842+
let type_ids = vec![value.geometry_type_id(); value.len()].into();
852843
let offsets = ScalarBuffer::from_iter(0..value.len() as i32);
853844
let data_type = GeometryType::new(metadata).with_coord_type(coord_type);
854845
let mut children = empty_children(coord_type);

‎rust/geoarrow-array/src/array/geometrycollection.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
55
use arrow_buffer::{NullBuffer, OffsetBuffer};
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, GeometryCollectionType, Metadata};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, GeometryCollectionType, Metadata};
910

1011
use crate::array::{GenericWkbArray, MixedGeometryArray};
1112
use crate::builder::GeometryCollectionBuilder;
@@ -271,6 +272,14 @@ impl PartialEq for GeometryCollectionArray {
271272
}
272273
}
273274

275+
impl GeometryTypeId for GeometryCollectionArray {
276+
const GEOMETRY_TYPE_OFFSET: i8 = 7;
277+
278+
fn dimension(&self) -> Dimension {
279+
self.data_type.dimension()
280+
}
281+
}
282+
274283
#[cfg(test)]
275284
mod test {
276285
use geoarrow_schema::{CoordType, Dimension};

‎rust/geoarrow-array/src/array/linestring.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
55
use arrow_buffer::{NullBuffer, OffsetBuffer};
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, LineStringType, Metadata};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, LineStringType, Metadata};
910

1011
use crate::array::{CoordBuffer, GenericWkbArray};
1112
use crate::builder::LineStringBuilder;
@@ -310,6 +311,14 @@ impl PartialEq for LineStringArray {
310311
}
311312
}
312313

314+
impl GeometryTypeId for LineStringArray {
315+
const GEOMETRY_TYPE_OFFSET: i8 = 2;
316+
317+
fn dimension(&self) -> Dimension {
318+
self.data_type.dimension()
319+
}
320+
}
321+
313322
#[cfg(test)]
314323
mod test {
315324
use geo_traits::to_geo::ToGeoLineString;

‎rust/geoarrow-array/src/array/mixed.rs‎

Lines changed: 30 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use arrow_array::{Array, ArrayRef, UnionArray};
66
use arrow_buffer::ScalarBuffer;
77
use arrow_schema::{DataType, UnionMode};
88
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
9+
use geoarrow_schema::type_id::GeometryTypeId;
910
use geoarrow_schema::{
1011
CoordType, Dimension, GeoArrowType, GeometryCollectionType, LineStringType,
1112
MultiLineStringType, MultiPointType, MultiPolygonType, PointType, PolygonType,
@@ -232,7 +233,7 @@ impl MixedGeometryArray {
232233
// If the array has been sliced, check a point type id still exists
233234
if self.is_sliced() {
234235
for t in self.type_ids.iter() {
235-
if *t % 10 == 1 {
236+
if *t % 10 == PointType::GEOMETRY_TYPE_OFFSET {
236237
return true;
237238
}
238239
}
@@ -251,7 +252,7 @@ impl MixedGeometryArray {
251252
// If the array has been sliced, check a point type id still exists
252253
if self.is_sliced() {
253254
for t in self.type_ids.iter() {
254-
if *t % 10 == 2 {
255+
if *t % 10 == LineStringType::GEOMETRY_TYPE_OFFSET {
255256
return true;
256257
}
257258
}
@@ -270,7 +271,7 @@ impl MixedGeometryArray {
270271
// If the array has been sliced, check a point type id still exists
271272
if self.is_sliced() {
272273
for t in self.type_ids.iter() {
273-
if *t % 10 == 3 {
274+
if *t % 10 == PolygonType::GEOMETRY_TYPE_OFFSET {
274275
return true;
275276
}
276277
}
@@ -289,7 +290,7 @@ impl MixedGeometryArray {
289290
// If the array has been sliced, check a point type id still exists
290291
if self.is_sliced() {
291292
for t in self.type_ids.iter() {
292-
if *t % 10 == 4 {
293+
if *t % 10 == MultiPointType::GEOMETRY_TYPE_OFFSET {
293294
return true;
294295
}
295296
}
@@ -308,7 +309,7 @@ impl MixedGeometryArray {
308309
// If the array has been sliced, check a point type id still exists
309310
if self.is_sliced() {
310311
for t in self.type_ids.iter() {
311-
if *t % 10 == 5 {
312+
if *t % 10 == MultiLineStringType::GEOMETRY_TYPE_OFFSET {
312313
return true;
313314
}
314315
}
@@ -327,7 +328,7 @@ impl MixedGeometryArray {
327328
// If the array has been sliced, check a point type id still exists
328329
if self.is_sliced() {
329330
for t in self.type_ids.iter() {
330-
if *t % 10 == 6 {
331+
if *t % 10 == MultiPolygonType::GEOMETRY_TYPE_OFFSET {
331332
return true;
332333
}
333334
}
@@ -492,15 +493,25 @@ impl MixedGeometryArray {
492493

493494
let expect_msg = "native geometry value access should never error";
494495
match type_id % 10 {
495-
1 => Geometry::Point(self.points.value(offset).expect(expect_msg)),
496-
2 => Geometry::LineString(self.line_strings.value(offset).expect(expect_msg)),
497-
3 => Geometry::Polygon(self.polygons.value(offset).expect(expect_msg)),
498-
4 => Geometry::MultiPoint(self.multi_points.value(offset).expect(expect_msg)),
499-
5 => {
496+
PointType::GEOMETRY_TYPE_OFFSET => {
497+
Geometry::Point(self.points.value(offset).expect(expect_msg))
498+
}
499+
LineStringType::GEOMETRY_TYPE_OFFSET => {
500+
Geometry::LineString(self.line_strings.value(offset).expect(expect_msg))
501+
}
502+
PolygonType::GEOMETRY_TYPE_OFFSET => {
503+
Geometry::Polygon(self.polygons.value(offset).expect(expect_msg))
504+
}
505+
MultiPointType::GEOMETRY_TYPE_OFFSET => {
506+
Geometry::MultiPoint(self.multi_points.value(offset).expect(expect_msg))
507+
}
508+
MultiLineStringType::GEOMETRY_TYPE_OFFSET => {
500509
Geometry::MultiLineString(self.multi_line_strings.value(offset).expect(expect_msg))
501510
}
502-
6 => Geometry::MultiPolygon(self.multi_polygons.value(offset).expect(expect_msg)),
503-
7 => {
511+
MultiPolygonType::GEOMETRY_TYPE_OFFSET => {
512+
Geometry::MultiPolygon(self.multi_polygons.value(offset).expect(expect_msg))
513+
}
514+
GeometryCollectionType::GEOMETRY_TYPE_OFFSET => {
504515
panic!("nested geometry collections not supported in GeoArrow")
505516
}
506517
_ => unreachable!("unknown type_id {}", type_id),
@@ -572,7 +583,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
572583
}
573584

574585
match type_id % 10 {
575-
1 => {
586+
PointType::GEOMETRY_TYPE_OFFSET => {
576587
points = Some(
577588
(
578589
value.child(type_id).as_ref(),
@@ -583,7 +594,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
583594
.unwrap(),
584595
);
585596
}
586-
2 => {
597+
LineStringType::GEOMETRY_TYPE_OFFSET => {
587598
line_strings = Some(
588599
(
589600
value.child(type_id).as_ref(),
@@ -594,7 +605,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
594605
.unwrap(),
595606
);
596607
}
597-
3 => {
608+
PolygonType::GEOMETRY_TYPE_OFFSET => {
598609
polygons = Some(
599610
(
600611
value.child(type_id).as_ref(),
@@ -605,7 +616,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
605616
.unwrap(),
606617
);
607618
}
608-
4 => {
619+
MultiPointType::GEOMETRY_TYPE_OFFSET => {
609620
multi_points = Some(
610621
(
611622
value.child(type_id).as_ref(),
@@ -616,7 +627,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
616627
.unwrap(),
617628
);
618629
}
619-
5 => {
630+
MultiLineStringType::GEOMETRY_TYPE_OFFSET => {
620631
multi_line_strings = Some(
621632
(
622633
value.child(type_id).as_ref(),
@@ -627,7 +638,7 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
627638
.unwrap(),
628639
);
629640
}
630-
6 => {
641+
MultiPolygonType::GEOMETRY_TYPE_OFFSET => {
631642
multi_polygons = Some(
632643
(
633644
value.child(type_id).as_ref(),

‎rust/geoarrow-array/src/array/multilinestring.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
55
use arrow_buffer::{NullBuffer, OffsetBuffer};
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, Metadata, MultiLineStringType};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, Metadata, MultiLineStringType};
910

1011
use crate::array::{CoordBuffer, GenericWkbArray, LineStringArray};
1112
use crate::builder::MultiLineStringBuilder;
@@ -393,6 +394,14 @@ impl PartialEq for MultiLineStringArray {
393394
}
394395
}
395396

397+
impl GeometryTypeId for MultiLineStringArray {
398+
const GEOMETRY_TYPE_OFFSET: i8 = 5;
399+
400+
fn dimension(&self) -> Dimension {
401+
self.data_type.dimension()
402+
}
403+
}
404+
396405
#[cfg(test)]
397406
mod test {
398407
use geo_traits::to_geo::ToGeoMultiLineString;

‎rust/geoarrow-array/src/array/multipoint.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
55
use arrow_buffer::{NullBuffer, OffsetBuffer};
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, Metadata, MultiPointType};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, Metadata, MultiPointType};
910

1011
use crate::array::{CoordBuffer, GenericWkbArray, PointArray};
1112
use crate::builder::MultiPointBuilder;
@@ -330,6 +331,14 @@ impl PartialEq for MultiPointArray {
330331
}
331332
}
332333

334+
impl GeometryTypeId for MultiPointArray {
335+
const GEOMETRY_TYPE_OFFSET: i8 = 4;
336+
337+
fn dimension(&self) -> Dimension {
338+
self.data_type.dimension()
339+
}
340+
}
341+
333342
#[cfg(test)]
334343
mod test {
335344
use geo_traits::to_geo::ToGeoMultiPoint;

‎rust/geoarrow-array/src/array/multipolygon.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
55
use arrow_buffer::{NullBuffer, OffsetBuffer};
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, Metadata, MultiPolygonType};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, Metadata, MultiPolygonType};
910

1011
use crate::array::{CoordBuffer, GenericWkbArray, PolygonArray};
1112
use crate::builder::MultiPolygonBuilder;
@@ -448,6 +449,14 @@ impl PartialEq for MultiPolygonArray {
448449
}
449450
}
450451

452+
impl GeometryTypeId for MultiPolygonArray {
453+
const GEOMETRY_TYPE_OFFSET: i8 = 6;
454+
455+
fn dimension(&self) -> Dimension {
456+
self.data_type.dimension()
457+
}
458+
}
459+
451460
#[cfg(test)]
452461
mod test {
453462
use geo_traits::to_geo::ToGeoMultiPolygon;

‎rust/geoarrow-array/src/array/point.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ use arrow_array::{Array, ArrayRef, FixedSizeListArray, StructArray};
55
use arrow_buffer::NullBuffer;
66
use arrow_schema::{DataType, Field};
77
use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8-
use geoarrow_schema::{CoordType, GeoArrowType, Metadata, PointType};
8+
use geoarrow_schema::type_id::GeometryTypeId;
9+
use geoarrow_schema::{CoordType, Dimension, GeoArrowType, Metadata, PointType};
910

1011
use crate::array::{CoordBuffer, InterleavedCoordBuffer, SeparatedCoordBuffer};
1112
use crate::eq::point_eq;
@@ -291,6 +292,14 @@ impl PartialEq for PointArray {
291292
}
292293
}
293294

295+
impl GeometryTypeId for PointArray {
296+
const GEOMETRY_TYPE_OFFSET: i8 = 1;
297+
298+
fn dimension(&self) -> Dimension {
299+
self.data_type.dimension()
300+
}
301+
}
302+
294303
#[cfg(test)]
295304
mod test {
296305
use geo_traits::to_geo::ToGeoPoint;

0 commit comments

Comments
 (0)