write 3d geometries to wkb (#669)

For #662
geoarrow · Jul 21, 2024 · fbd495b · fbd495b
1 parent 3fc250d
commit fbd495b
Show file tree

Hide file tree

Showing 17 changed files with 271 additions and 127 deletions.
diff --git a/src/algorithm/native/type_id.rs b/src/algorithm/native/type_id.rs
@@ -1,9 +1,8 @@
 use crate::array::*;
-use crate::io::wkb::reader::WKBGeometryType;
 use crate::trait_::GeometryArrayAccessor;
 use crate::GeometryArrayTrait;
-use arrow_array::builder::Int8Builder;
-use arrow_array::{Int8Array, OffsetSizeTrait};
+use arrow::array::Int16Builder;
+use arrow_array::{Int16Array, OffsetSizeTrait};
 use std::collections::HashSet;
 
 /// Calculation of the geometry types within a GeometryArray
@@ -21,7 +20,7 @@ pub trait TypeIds {
     /// - MULTILINESTRING is 5
     /// - MULTIPOLYGON is 6
     /// - GEOMETRYCOLLECTION is 7
-    fn get_type_ids(&self) -> Int8Array;
+    fn get_type_ids(&self) -> Int16Array;
 
     /// Return the unique geometry types stored in this array
     ///
@@ -36,16 +35,16 @@ pub trait TypeIds {
     /// - MULTILINESTRING is 5
     /// - MULTIPOLYGON is 6
     /// - GEOMETRYCOLLECTION is 7
-    fn get_unique_type_ids(&self) -> HashSet<i8>;
+    fn get_unique_type_ids(&self) -> HashSet<i16>;
 }
 
 impl TypeIds for PointArray<2> {
-    fn get_type_ids(&self) -> Int8Array {
-        let values = vec![0i8; self.len()];
-        Int8Array::new(values.into(), self.nulls().cloned())
+    fn get_type_ids(&self) -> Int16Array {
+        let values = vec![0i16; self.len()];
+        Int16Array::new(values.into(), self.nulls().cloned())
     }
 
-    fn get_unique_type_ids(&self) -> HashSet<i8> {
+    fn get_unique_type_ids(&self) -> HashSet<i16> {
         let mut values = HashSet::with_capacity(1);
         values.insert(0);
         values
@@ -55,12 +54,12 @@ impl TypeIds for PointArray<2> {
 macro_rules! constant_impl {
     ($type:ty, $value:expr) => {
         impl<O: OffsetSizeTrait> TypeIds for $type {
-            fn get_type_ids(&self) -> Int8Array {
+            fn get_type_ids(&self) -> Int16Array {
                 let values = vec![$value; self.len()];
-                Int8Array::new(values.into(), self.nulls().cloned())
+                Int16Array::new(values.into(), self.nulls().cloned())
             }
 
-            fn get_unique_type_ids(&self) -> HashSet<i8> {
+            fn get_unique_type_ids(&self) -> HashSet<i16> {
                 let mut values = HashSet::with_capacity(1);
                 values.insert($value);
                 values
@@ -76,10 +75,10 @@ constant_impl!(MultiLineStringArray<O, 2>, 5);
 constant_impl!(MultiPolygonArray<O, 2>, 6);
 
 impl<O: OffsetSizeTrait> TypeIds for MixedGeometryArray<O, 2> {
-    fn get_type_ids(&self) -> Int8Array {
+    fn get_type_ids(&self) -> Int16Array {
         use crate::scalar::Geometry::*;
 
-        let mut output_array = Int8Builder::with_capacity(self.len());
+        let mut output_array = Int16Builder::with_capacity(self.len());
         self.iter().for_each(|maybe_g| {
             output_array.append_option(maybe_g.map(|g| match g {
                 Point(_) => 0,
@@ -95,7 +94,7 @@ impl<O: OffsetSizeTrait> TypeIds for MixedGeometryArray<O, 2> {
         output_array.finish()
     }
 
-    fn get_unique_type_ids(&self) -> HashSet<i8> {
+    fn get_unique_type_ids(&self) -> HashSet<i16> {
         use crate::scalar::Geometry::*;
 
         let mut values = HashSet::new();
@@ -118,41 +117,24 @@ impl<O: OffsetSizeTrait> TypeIds for MixedGeometryArray<O, 2> {
 }
 
 impl<O: OffsetSizeTrait> TypeIds for WKBArray<O> {
-    fn get_type_ids(&self) -> Int8Array {
-        use WKBGeometryType::*;
-
-        let mut output_array = Int8Builder::with_capacity(self.len());
+    fn get_type_ids(&self) -> Int16Array {
+        let mut output_array = Int16Builder::with_capacity(self.len());
 
         self.iter().for_each(|maybe_wkb| {
-            output_array.append_option(maybe_wkb.map(|wkb| match wkb.get_wkb_geometry_type() {
-                Point => 0,
-                LineString => 1,
-                Polygon => 3,
-                MultiPoint => 4,
-                MultiLineString => 5,
-                MultiPolygon => 6,
-                GeometryCollection => 7,
+            output_array.append_option(maybe_wkb.map(|wkb| {
+                let type_id = u32::from(wkb.get_wkb_geometry_type());
+                type_id.try_into().unwrap()
             }))
         });
 
         output_array.finish()
     }
 
-    fn get_unique_type_ids(&self) -> HashSet<i8> {
-        use WKBGeometryType::*;
-
+    fn get_unique_type_ids(&self) -> HashSet<i16> {
         let mut values = HashSet::new();
         self.iter().flatten().for_each(|wkb| {
-            let type_id = match wkb.get_wkb_geometry_type() {
-                Point => 0,
-                LineString => 1,
-                Polygon => 3,
-                MultiPoint => 4,
-                MultiLineString => 5,
-                MultiPolygon => 6,
-                GeometryCollection => 7,
-            };
-            values.insert(type_id);
+            let type_id = u32::from(wkb.get_wkb_geometry_type());
+            values.insert(type_id.try_into().unwrap());
         });
 
         values

diff --git a/src/array/binary/builder.rs b/src/array/binary/builder.rs
@@ -9,10 +9,9 @@ use crate::geo_traits::{
 };
 use crate::io::wkb::writer::{
     geometry_collection_wkb_size, line_string_wkb_size, multi_line_string_wkb_size,
-    multi_point_wkb_size, multi_polygon_wkb_size, polygon_wkb_size,
+    multi_point_wkb_size, multi_polygon_wkb_size, point_wkb_size, polygon_wkb_size,
     write_geometry_collection_as_wkb, write_line_string_as_wkb, write_multi_line_string_as_wkb,
     write_multi_point_as_wkb, write_multi_polygon_as_wkb, write_point_as_wkb, write_polygon_as_wkb,
-    POINT_WKB_SIZE,
 };
 use arrow_array::builder::GenericBinaryBuilder;
 use arrow_array::OffsetSizeTrait;
@@ -82,7 +81,7 @@ impl<O: OffsetSizeTrait> WKBBuilder<O> {
     pub fn push_point(&mut self, geom: Option<&impl PointTrait<T = f64>>) {
         if let Some(geom) = geom {
             // TODO: figure out how to write directly to the underlying vec without a copy
-            let mut buf = Vec::with_capacity(POINT_WKB_SIZE);
+            let mut buf = Vec::with_capacity(point_wkb_size(geom.dim()));
             write_point_as_wkb(&mut buf, geom).unwrap();
             self.0.append_value(&buf)
         } else {

diff --git a/src/array/binary/capacity.rs b/src/array/binary/capacity.rs
@@ -8,7 +8,7 @@ use crate::geo_traits::{
 };
 use crate::io::wkb::writer::{
     geometry_collection_wkb_size, line_string_wkb_size, multi_line_string_wkb_size,
-    multi_point_wkb_size, multi_polygon_wkb_size, polygon_wkb_size, POINT_WKB_SIZE,
+    multi_point_wkb_size, multi_polygon_wkb_size, point_wkb_size, polygon_wkb_size,
 };
 
 /// A counter for the buffer sizes of a [`WKBArray`][crate::array::WKBArray].
@@ -49,9 +49,9 @@ impl WKBCapacity {
 
     /// Add a Point to this capacity counter.
     #[inline]
-    pub fn add_point(&mut self, is_valid: bool) {
-        if is_valid {
-            self.buffer_capacity += POINT_WKB_SIZE;
+    pub fn add_point<'a>(&mut self, point: Option<&'a (impl PointTrait + 'a)>) {
+        if let Some(point) = point {
+            self.buffer_capacity += point_wkb_size(point.dim());
         }
         self.offsets_capacity += 1;
     }
@@ -112,7 +112,7 @@ impl WKBCapacity {
     pub fn add_geometry<'a>(&mut self, geom: Option<&'a (impl GeometryTrait + 'a)>) {
         if let Some(geom) = geom {
             match geom.as_type() {
-                crate::geo_traits::GeometryType::Point(_) => self.add_point(true),
+                crate::geo_traits::GeometryType::Point(g) => self.add_point(Some(g)),
                 crate::geo_traits::GeometryType::LineString(g) => self.add_line_string(Some(g)),
                 crate::geo_traits::GeometryType::Polygon(g) => self.add_polygon(Some(g)),
                 crate::geo_traits::GeometryType::MultiPoint(p) => self.add_multi_point(Some(p)),
@@ -148,7 +148,7 @@ impl WKBCapacity {
     ) -> Self {
         let mut counter = Self::new_empty();
         for maybe_geom in geoms.into_iter() {
-            counter.add_point(maybe_geom.is_some());
+            counter.add_point(maybe_geom);
         }
         counter
     }

diff --git a/src/io/wkb/common.rs b/src/io/wkb/common.rs
@@ -0,0 +1,20 @@
+use num_enum::{IntoPrimitive, TryFromPrimitive};
+
+#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive, IntoPrimitive)]
+#[repr(u32)]
+pub enum WKBType {
+    Point = 1,
+    LineString = 2,
+    Polygon = 3,
+    MultiPoint = 4,
+    MultiLineString = 5,
+    MultiPolygon = 6,
+    GeometryCollection = 7,
+    PointZ = 1001,
+    LineStringZ = 1002,
+    PolygonZ = 1003,
+    MultiPointZ = 1004,
+    MultiLineStringZ = 1005,
+    MultiPolygonZ = 1006,
+    GeometryCollectionZ = 1007,
+}
diff --git a/src/io/wkb/mod.rs b/src/io/wkb/mod.rs
@@ -1,6 +1,7 @@
 //! An optimized implementation of reading and writing ISO-flavored WKB-encoded geometries.
 
 mod api;
+pub(crate) mod common;
 pub(crate) mod reader;
 pub(crate) mod writer;
 

diff --git a/src/io/wkb/reader/geometry.rs b/src/io/wkb/reader/geometry.rs
@@ -8,11 +8,12 @@ use crate::geo_traits::{
     GeometryCollectionTrait, GeometryTrait, LineStringTrait, MultiLineStringTrait, MultiPointTrait,
     MultiPolygonTrait, PointTrait, PolygonTrait,
 };
+use crate::io::wkb::common::WKBType;
 use crate::io::wkb::reader::geometry_collection::WKBGeometryCollection;
 use crate::io::wkb::reader::rect::WKBRect;
 use crate::io::wkb::reader::{
-    WKBGeometryType, WKBLineString, WKBMaybeMultiLineString, WKBMaybeMultiPoint,
-    WKBMaybeMultiPolygon, WKBMultiLineString, WKBMultiPoint, WKBMultiPolygon, WKBPoint, WKBPolygon,
+    WKBLineString, WKBMaybeMultiLineString, WKBMaybeMultiPoint, WKBMaybeMultiPolygon,
+    WKBMultiLineString, WKBMultiPoint, WKBMultiPolygon, WKBPoint, WKBPolygon,
 };
 use crate::scalar::WKB;
 
@@ -21,70 +22,74 @@ impl<'a, O: OffsetSizeTrait> WKB<'a, O> {
         let buf = self.arr.value(self.geom_index);
         let mut reader = Cursor::new(buf);
         let byte_order = reader.read_u8().unwrap();
-        let geometry_type = match byte_order {
+        let geometry_type_u32 = match byte_order {
             0 => reader.read_u32::<BigEndian>().unwrap(),
             1 => reader.read_u32::<LittleEndian>().unwrap(),
             _ => panic!("Unexpected byte order."),
         };
+        let geometry_type = WKBType::try_from(geometry_type_u32).unwrap();
 
         match geometry_type {
-            1 => WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XY)),
-            2 => WKBGeometry::LineString(WKBLineString::new(
+            WKBType::Point => {
+                WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XY))
+            }
+            WKBType::LineString => WKBGeometry::LineString(WKBLineString::new(
                 buf,
                 byte_order.into(),
                 0,
                 Dimension::XY,
             )),
-            3 => WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, Dimension::XY)),
-            4 => WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), Dimension::XY)),
-            5 => WKBGeometry::MultiLineString(WKBMultiLineString::new(
-                buf,
-                byte_order.into(),
-                Dimension::XY,
-            )),
-            6 => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
+            WKBType::Polygon => {
+                WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, Dimension::XY))
+            }
+            WKBType::MultiPoint => {
+                WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), Dimension::XY))
+            }
+            WKBType::MultiLineString => WKBGeometry::MultiLineString(WKBMultiLineString::new(
                 buf,
                 byte_order.into(),
                 Dimension::XY,
             )),
-            7 => WKBGeometry::GeometryCollection(WKBGeometryCollection::new(
+            WKBType::MultiPolygon => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
                 buf,
                 byte_order.into(),
                 Dimension::XY,
             )),
-            1001 => WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XYZ)),
-            1002 => WKBGeometry::LineString(WKBLineString::new(
+            WKBType::GeometryCollection => WKBGeometry::GeometryCollection(
+                WKBGeometryCollection::new(buf, byte_order.into(), Dimension::XY),
+            ),
+            WKBType::PointZ => {
+                WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XYZ))
+            }
+            WKBType::LineStringZ => WKBGeometry::LineString(WKBLineString::new(
                 buf,
                 byte_order.into(),
                 0,
                 Dimension::XYZ,
             )),
-            1003 => {
+            WKBType::PolygonZ => {
                 WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, Dimension::XYZ))
             }
-            1004 => {
+            WKBType::MultiPointZ => {
                 WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), Dimension::XYZ))
             }
-            1005 => WKBGeometry::MultiLineString(WKBMultiLineString::new(
-                buf,
-                byte_order.into(),
-                Dimension::XYZ,
-            )),
-            1006 => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
+            WKBType::MultiLineStringZ => WKBGeometry::MultiLineString(WKBMultiLineString::new(
                 buf,
                 byte_order.into(),
                 Dimension::XYZ,
             )),
-            1007 => WKBGeometry::GeometryCollection(WKBGeometryCollection::new(
+            WKBType::MultiPolygonZ => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
                 buf,
                 byte_order.into(),
                 Dimension::XYZ,
             )),
-            _ => panic!("Unsupported geometry type"),
+            WKBType::GeometryCollectionZ => WKBGeometry::GeometryCollection(
+                WKBGeometryCollection::new(buf, byte_order.into(), Dimension::XYZ),
+            ),
         }
     }
 
-    pub fn get_wkb_geometry_type(&'a self) -> WKBGeometryType {
+    pub fn get_wkb_geometry_type(&'a self) -> WKBType {
         let buf = self.arr.value(self.geom_index);
         let mut reader = Cursor::new(buf);
         let byte_order = reader.read_u8().unwrap();

diff --git a/src/io/wkb/reader/mod.rs b/src/io/wkb/reader/mod.rs
@@ -31,4 +31,3 @@ pub use multipoint::WKBMultiPoint;
 pub use multipolygon::WKBMultiPolygon;
 pub use point::WKBPoint;
 pub use polygon::WKBPolygon;
-pub use r#type::WKBGeometryType;
diff --git a/src/io/wkb/reader/type.rs b/src/io/wkb/reader/type.rs
@@ -7,6 +7,7 @@ use num_enum::TryFromPrimitive;
 use crate::array::CoordType;
 use crate::datatypes::{Dimension, GeoDataType};
 use crate::error::{GeoArrowError, Result};
+use crate::io::wkb::common::WKBType;
 use crate::scalar::WKB;
 
 #[derive(Debug, Eq, PartialEq, TryFromPrimitive)]
@@ -174,13 +175,14 @@ pub(crate) fn infer_geometry_type<'a, O: OffsetSizeTrait>(
     let mut available_type = AvailableTypes::new();
     for geom in geoms {
         match geom.get_wkb_geometry_type() {
-            WKBGeometryType::Point => available_type.add_point(),
-            WKBGeometryType::LineString => available_type.add_line_string(),
-            WKBGeometryType::Polygon => available_type.add_polygon(),
-            WKBGeometryType::MultiPoint => available_type.add_multi_point(),
-            WKBGeometryType::MultiLineString => available_type.add_multi_line_string(),
-            WKBGeometryType::MultiPolygon => available_type.add_multi_polygon(),
-            WKBGeometryType::GeometryCollection => available_type.add_geometry_collection(),
+            WKBType::Point => available_type.add_point(),
+            WKBType::LineString => available_type.add_line_string(),
+            WKBType::Polygon => available_type.add_polygon(),
+            WKBType::MultiPoint => available_type.add_multi_point(),
+            WKBType::MultiLineString => available_type.add_multi_line_string(),
+            WKBType::MultiPolygon => available_type.add_multi_polygon(),
+            WKBType::GeometryCollection => available_type.add_geometry_collection(),
+            _ => todo!("3d support"),
         }
     }
     available_type.resolve_type(large_type, coord_type)