FlatGeobuf 3d support (#728)

This is failing on unimplemented 3d downcasting

### Change list

- Use geo-traits implementation from
#779
- Make a couple fixes in the geo-traits.
- Offsets were _coordinate_ offsets and need to be multiplied by 2 to
get the correct x and y values
    - Lengths are `end - start`
- I was computing the number of linestrings in multi line strings
incorrectly.
- Add more tests on the Python side.

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)