FlatGeobuf 3d support

js/tests/js/flatgeobuf.test.ts

@@ -18,6 +18,6 @@ it("read FlatGeobuf", () => {
   const geometryFieldMetadata = geometryField.metadata;
   console.log(geometryFieldMetadata);
   expect(geometryFieldMetadata.get("ARROW:extension:name")).toStrictEqual(
-    "geoarrow.polygon"
+    "geoarrow.multipolygon"
   );
 });

python/tests/io/test_flatgeobuf.py

@@ -1,9 +1,12 @@
 from io import BytesIO
 
-from geoarrow.rust.io import read_flatgeobuf, write_flatgeobuf
-from geoarrow.rust.core import to_geopandas
+import geodatasets
 import geopandas as gpd
+import pyarrow as pa
 import pytest
+import shapely
+from geoarrow.rust.core import from_geopandas, geometry_col, to_geopandas
+from geoarrow.rust.io import read_flatgeobuf, write_flatgeobuf
 from geopandas.testing import assert_geodataframe_equal
 
 from tests.utils import FIXTURES_DIR
@@ -13,28 +16,88 @@ def test_read_flatgeobuf():
     path = FIXTURES_DIR / "flatgeobuf" / "countries.fgb"
     table = read_flatgeobuf(path)
     assert len(table) == 179
-    # assert isinstance(gars.geometry_col(table), gars.ChunkedMultiPolygonArray)
+    # hacky
+    assert "MultiPolygon" in geometry_col(table).type.__repr__()
 
 
 def test_read_flatgeobuf_file_object():
     path = FIXTURES_DIR / "flatgeobuf" / "countries.fgb"
     with open(path, "rb") as f:
         table = read_flatgeobuf(f)
     assert len(table) == 179
-    # assert isinstance(gars.geometry_col(table), gars.ChunkedMultiPolygonArray)
+    # hacky
+    assert "MultiPolygon" in geometry_col(table).type.__repr__()
 
 
 def test_round_trip_flatgeobuf():
     path = FIXTURES_DIR / "flatgeobuf" / "countries.fgb"
     table = read_flatgeobuf(path)
 
     buf = BytesIO()
-    write_flatgeobuf(table, buf)
+    write_flatgeobuf(table, buf, write_index=False)
     buf.seek(0)
     table_back = read_flatgeobuf(buf)
     assert table == table_back  # type: ignore
 
 
+def test_round_trip_polygon():
+    geom = shapely.geometry.shape(
+        {
+            "type": "Polygon",
+            "coordinates": [
+                [
+                    [-118.4765625, 33.92578125],
+                    [-118.125, 33.92578125],
+                    [-118.125, 34.1015625],
+                    [-118.4765625, 34.1015625],
+                    [-118.4765625, 33.92578125],
+                ],
+                [
+                    [-118.24447631835938, 34.0521240234375],
+                    [-118.24310302734375, 34.0521240234375],
+                    [-118.24310302734375, 34.053497314453125],
+                    [-118.24447631835938, 34.053497314453125],
+                    [-118.24447631835938, 34.0521240234375],
+                ],
+            ],
+        }
+    )
+    polys = [geom] * 3
+    gdf = gpd.GeoDataFrame({"col1": ["a", "b", "c"]}, geometry=polys, crs="EPSG:4326")
+    table = from_geopandas(gdf)
+
+    buf = BytesIO()
+    write_flatgeobuf(table, buf, write_index=False)
+    buf.seek(0)
+    table_back = read_flatgeobuf(buf)
+    assert pa.table(table) == pa.table(table_back)
+
+
+def test_round_trip_3d_points():
+    points = shapely.points([1, 2, 3], [4, 5, 6], [7, 8, 9])
+    gdf = gpd.GeoDataFrame({"col1": ["a", "b", "c"]}, geometry=points, crs="EPSG:4326")
+    table = from_geopandas(gdf)
+
+    buf = BytesIO()
+    write_flatgeobuf(table, buf, write_index=False)
+    buf.seek(0)
+    table_back = read_flatgeobuf(buf)
+
+    assert pa.table(table) == pa.table(table_back)
+
+
+def test_round_trip_multilinestring():
+    gdf = gpd.read_file(geodatasets.get_path("eea.large_rivers"))
+    table = from_geopandas(gdf)
+
+    buf = BytesIO()
+    write_flatgeobuf(table, buf, write_index=False)
+    buf.seek(0)
+    table_back = read_flatgeobuf(buf)
+
+    assert pa.table(table) == pa.table(table_back)
+
+
 @pytest.mark.xfail(reason="fix propagate CRS")
 def test_matches_pyogrio():
     path = FIXTURES_DIR / "flatgeobuf" / "countries.fgb"

src/io/flatgeobuf/reader/async.rs

@@ -1,12 +1,14 @@
 use std::sync::Arc;
 
 use flatgeobuf::{GeometryType, HttpFgbReader};
+use geozero::{FeatureProcessor, FeatureProperties};
 use http_range_client::AsyncBufferedHttpRangeClient;
 use object_store::path::Path;
 use object_store::ObjectStore;
 
 use crate::algorithm::native::DowncastTable;
 use crate::array::*;
+use crate::datatypes::Dimension;
 use crate::error::{GeoArrowError, Result};
 use crate::io::flatgeobuf::reader::common::{infer_schema, FlatGeobufReaderOptions};
 use crate::io::flatgeobuf::reader::object_store_reader::ObjectStoreWrapper;
@@ -31,11 +33,12 @@ pub async fn read_flatgeobuf_async(
     let reader = HttpFgbReader::new(async_client).await.unwrap();
 
     let header = reader.header();
-    if header.has_m() | header.has_t() | header.has_tm() | header.has_z() {
+    if header.has_m() | header.has_t() | header.has_tm() {
         return Err(GeoArrowError::General(
-            "Only XY dimensions are supported".to_string(),
+            "Only XY and XYZ dimensions are supported".to_string(),
         ));
     }
+    let has_z = header.has_z();
 
     let schema = infer_schema(header);
     let geometry_type = header.geometry_type();
@@ -58,41 +61,89 @@ pub async fn read_flatgeobuf_async(
         Default::default(),
     );
 
-    match geometry_type {
-        GeometryType::Point => {
-            let mut builder = GeoTableBuilder::<PointBuilder<2>>::new_with_options(options);
+    macro_rules! impl_read {
+        ($builder:ty, $geom_type:ty, $dim:expr) => {{
+            let mut builder = GeoTableBuilder::<$builder>::new_with_options(options);
+            while let Some(feature) = selection.next().await? {
+                feature.process_properties(&mut builder)?;
+                builder.properties_end()?;
+
+                let geom: Option<super::core::Geometry<'_>> = feature
+                    .geometry()
+                    .map(|g| <$geom_type>::new(g, $dim))
+                    .map(|g| g.into());
+                builder.push_geometry(geom.as_ref())?;
+
+                builder.feature_end(0)?;
+            }
             selection.process_features(&mut builder).await?;
             builder.finish()
+        }};
+    }
+
+    match (geometry_type, has_z) {
+        (GeometryType::Point, false) => {
+            impl_read!(PointBuilder<2>, super::core::Point, Dimension::XY)
         }
-        GeometryType::LineString => {
-            let mut builder = GeoTableBuilder::<LineStringBuilder<2>>::new_with_options(options);
-            selection.process_features(&mut builder).await?;
-            builder.finish()
+        (GeometryType::LineString, false) => {
+            impl_read!(LineStringBuilder<2>, super::core::LineString, Dimension::XY)
         }
-        GeometryType::Polygon => {
-            let mut builder = GeoTableBuilder::<PolygonBuilder<2>>::new_with_options(options);
-            selection.process_features(&mut builder).await?;
-            builder.finish()
+        (GeometryType::Polygon, false) => {
+            impl_read!(PolygonBuilder<2>, super::core::Polygon, Dimension::XY)
         }
-        GeometryType::MultiPoint => {
-            let mut builder = GeoTableBuilder::<MultiPointBuilder<2>>::new_with_options(options);
-            selection.process_features(&mut builder).await?;
-            builder.finish()
+        (GeometryType::MultiPoint, false) => {
+            impl_read!(MultiPointBuilder<2>, super::core::MultiPoint, Dimension::XY)
         }
-        GeometryType::MultiLineString => {
+        (GeometryType::MultiLineString, false) => impl_read!(
+            MultiLineStringBuilder<2>,
+            super::core::MultiLineString,
+            Dimension::XY
+        ),
+        (GeometryType::MultiPolygon, false) => impl_read!(
+            MultiPolygonBuilder<2>,
+            super::core::MultiPolygon,
+            Dimension::XY
+        ),
+        (GeometryType::Unknown, false) => {
             let mut builder =
-                GeoTableBuilder::<MultiLineStringBuilder<2>>::new_with_options(options);
+                GeoTableBuilder::<MixedGeometryStreamBuilder<2>>::new_with_options(options);
             selection.process_features(&mut builder).await?;
-            builder.finish()
+            let table = builder.finish()?;
+            table.downcast(true)
         }
-        GeometryType::MultiPolygon => {
-            let mut builder = GeoTableBuilder::<MultiPolygonBuilder<2>>::new_with_options(options);
-            selection.process_features(&mut builder).await?;
-            builder.finish()
+        (GeometryType::Point, true) => {
+            impl_read!(PointBuilder<3>, super::core::Point, Dimension::XYZ)
         }
-        GeometryType::Unknown => {
+        (GeometryType::LineString, true) => {
+            impl_read!(
+                LineStringBuilder<3>,
+                super::core::LineString,
+                Dimension::XYZ
+            )
+        }
+        (GeometryType::Polygon, true) => {
+            impl_read!(PolygonBuilder<3>, super::core::Polygon, Dimension::XYZ)
+        }
+        (GeometryType::MultiPoint, true) => {
+            impl_read!(
+                MultiPointBuilder<3>,
+                super::core::MultiPoint,
+                Dimension::XYZ
+            )
+        }
+        (GeometryType::MultiLineString, true) => impl_read!(
+            MultiLineStringBuilder<3>,
+            super::core::MultiLineString,
+            Dimension::XYZ
+        ),
+        (GeometryType::MultiPolygon, true) => impl_read!(
+            MultiPolygonBuilder<3>,
+            super::core::MultiPolygon,
+            Dimension::XYZ
+        ),
+        (GeometryType::Unknown, true) => {
             let mut builder =
-                GeoTableBuilder::<MixedGeometryStreamBuilder<2>>::new_with_options(options);
+                GeoTableBuilder::<MixedGeometryStreamBuilder<3>>::new_with_options(options);
             selection.process_features(&mut builder).await?;
             let table = builder.finish()?;
             table.downcast(true)