refactor: some of the geoparquet code (#768)

gadomski · web-flow · commit a591cf2a0233 · 2025-06-04T14:50:01.000Z
diff --git a/Cargo.toml b/Cargo.toml
@@ -56,9 +56,9 @@ futures-util = "0.3.31"
 geo = "0.30.0"
 geo-traits = "0.3.0"
 geo-types = "0.7.16"
-geoarrow-array = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "d27500849c6cee019535d6749991d1fd122baecf" }
-geoparquet = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "d27500849c6cee019535d6749991d1fd122baecf" }
-geoarrow-schema = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "d27500849c6cee019535d6749991d1fd122baecf" }
+geoarrow-array = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "1d73a8c8f739ac2a9f4cbac928283c19b74db463" }
+geoparquet = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "1d73a8c8f739ac2a9f4cbac928283c19b74db463" }
+geoarrow-schema = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "1d73a8c8f739ac2a9f4cbac928283c19b74db463" }
 geojson = "0.24.1"
 getrandom = { version = "0.3.3", features = ["wasm_js"] }
 http = "1.1"
diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs
@@ -58,6 +58,11 @@ pub enum Error {
     #[error("no \"{0}\" field in the JSON object")]
     MissingField(&'static str),
 
+    /// No geoparquet metadata in a stac-geoparquet file.
+    #[error("no geoparquet metadata")]
+    #[cfg(feature = "geoparquet")]
+    MissingGeoparquetMetadata,
+
     /// There are no items, when items are required.
     #[error("no items")]
     NoItems,
diff --git a/crates/core/src/geoarrow/mod.rs b/crates/core/src/geoarrow/mod.rs
@@ -12,7 +12,7 @@ use geoarrow_array::{
     array::{WkbArray, from_arrow_array},
     builder::GeometryBuilder,
 };
-use geoarrow_schema::{CoordType, GeoArrowType, GeometryType, Metadata};
+use geoarrow_schema::{GeoArrowType, GeometryType, Metadata};
 use serde_json::{Value, json};
 use std::{collections::HashMap, sync::Arc};
 
@@ -22,10 +22,10 @@ pub const VERSION_KEY: &str = "stac:geoparquet_version";
 /// The stac-geoparquet version.
 pub const VERSION: &str = "1.0.0";
 
-/// Other geoarrow geometry columns (other than "geometry")
-pub const OTHER_GEOMETRY_COLUMNS: [&str; 1] = ["proj:geometry"];
+/// Geometry columns.
+pub const GEOMETRY_COLUMNS: [&str; 2] = ["geometry", "proj:geometry"];
 
-/// Geoarrow datetime columns
+/// Datetime columns.
 pub const DATETIME_COLUMNS: [&str; 8] = [
     "datetime",
     "start_datetime",
@@ -37,7 +37,7 @@ pub const DATETIME_COLUMNS: [&str; 8] = [
     "unpublished",
 ];
 
-/// A geoarrow table.
+/// A **stac-geoarrow** table.
 ///
 /// `Table` existed in **geoarrow** v0.3 but was removed in v0.4. We preserve it
 /// here as a useful arrow-ish analog to [ItemCollection].
@@ -48,19 +48,6 @@ pub struct Table {
 }
 
 /// A builder for converting an [ItemCollection] to a [Table]
-///
-/// # Examples
-///
-/// ```
-/// use stac::geoarrow::TableBuilder;
-///
-/// let item = stac::read("examples/simple-item.json").unwrap();
-/// let builder = TableBuilder {
-///     item_collection: vec![item].into(),
-///     drop_invalid_attributes: false,
-/// };
-/// let table = builder.build().unwrap();
-/// ```
 #[derive(Debug)]
 pub struct TableBuilder {
     /// The item collection.
@@ -70,73 +57,60 @@ pub struct TableBuilder {
     ///
     /// If false, an invalid attribute will cause an error. If true, an invalid
     /// attribute will trigger a warning.
+    ///
+    /// Invalid attributes are values in `properties` that would conflict with a STAC-defined top-level key.
     pub drop_invalid_attributes: bool,
 }
 
 impl TableBuilder {
-    /// Builds a [Table]
+    /// Builds a [Table].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use stac::geoarrow::TableBuilder;
+    ///
+    /// let item = stac::read("examples/simple-item.json").unwrap();
+    /// let builder = TableBuilder {
+    ///     item_collection: vec![item].into(),
+    ///     drop_invalid_attributes: false,
+    /// };
+    /// let table = builder.build().unwrap();
+    /// ```
     pub fn build(self) -> Result<Table> {
         let mut values = Vec::with_capacity(self.item_collection.items.len());
-        let geometry_type = GeometryType::new(CoordType::Interleaved, Default::default()); // the choice of interleaved is arbitrary
-        let mut builder = GeometryBuilder::new(geometry_type.clone());
         let mut geometry_builders = HashMap::new();
-        for mut item in self.item_collection.items {
-            builder.push_geometry(
-                item.geometry
-                    .take()
-                    .and_then(|geometry| Geometry::try_from(geometry).ok())
-                    .as_ref(),
-            )?;
-            let flat_item = item.into_flat_item(self.drop_invalid_attributes)?;
-            let mut value = serde_json::to_value(flat_item)?;
+
+        for item in self.item_collection.items {
+            let mut value =
+                serde_json::to_value(item.into_flat_item(self.drop_invalid_attributes)?)?;
             {
                 let value = value
                     .as_object_mut()
                     .expect("a flat item should serialize to an object");
-                for key in OTHER_GEOMETRY_COLUMNS {
+                for key in GEOMETRY_COLUMNS {
                     if let Some(value) = value.remove(key) {
-                        let entry = geometry_builders
-                            .entry(key)
-                            .or_insert_with(|| GeometryBuilder::new(geometry_type.clone()));
+                        let entry = geometry_builders.entry(key).or_insert_with(|| {
+                            let geometry_type = GeometryType::new(Default::default());
+                            GeometryBuilder::new(geometry_type)
+                        });
                         let geometry =
                             geojson::Geometry::from_json_value(value).map_err(Box::new)?;
                         entry.push_geometry(Some(
                             &(Geometry::try_from(geometry).map_err(Box::new)?),
                         ))?;
                     }
                 }
-                let _ = value.remove("geometry");
                 if let Some(bbox) = value.remove("bbox") {
-                    let bbox = bbox
-                        .as_array()
-                        .expect("STAC items should always have a list as their bbox");
-                    if bbox.len() == 4 {
-                        let _ = value.insert("bbox".into(), json!({
-                        "xmin": bbox[0].as_number().expect("all bbox values should be a number"),
-                        "ymin": bbox[1].as_number().expect("all bbox values should be a number"),
-                        "xmax": bbox[2].as_number().expect("all bbox values should be a number"),
-                        "ymax": bbox[3].as_number().expect("all bbox values should be a number"),
-                    }));
-                    } else if bbox.len() == 6 {
-                        let _ = value.insert("bbox".into(), json!({
-                        "xmin": bbox[0].as_number().expect("all bbox values should be a number"),
-                        "ymin": bbox[1].as_number().expect("all bbox values should be a number"),
-                        "zmin": bbox[2].as_number().expect("all bbox values should be a number"),
-                        "xmax": bbox[3].as_number().expect("all bbox values should be a number"),
-                        "ymax": bbox[4].as_number().expect("all bbox values should be a number"),
-                        "zmax": bbox[5].as_number().expect("all bbox values should be a number"),
-                    }));
-                    } else {
-                        return Err(Error::InvalidBbox(
-                            bbox.iter().filter_map(|v| v.as_f64()).collect(),
-                        ));
-                    }
+                    let bbox = convert_bbox(bbox)?;
+                    let _ = value.insert("bbox".to_string(), bbox);
                 }
             }
             values.push(value);
         }
 
-        // Decode JSON
+        // Create a geometry-less record batch of our items.
+        // TODO do this in one pass: https://github.com/stac-utils/rustac/issues/767
         let schema = arrow_json::reader::infer_json_schema_from_iterator(values.iter().map(Ok))?;
         let mut schema_builder = SchemaBuilder::new();
         for field in schema.fields().iter() {
@@ -150,29 +124,23 @@ impl TableBuilder {
                 schema_builder.push(field.clone());
             }
         }
-        let mut metadata = schema.metadata;
-        let _ = metadata.insert(VERSION_KEY.to_string(), VERSION.into());
         let schema = Arc::new(schema_builder.finish());
         let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder()?;
         decoder.serialize(&values)?;
-
-        // Build the table
         let record_batch = decoder.flush()?.ok_or(Error::NoItems)?;
+
+        // Add the geometries back in.
         let mut schema_builder = SchemaBuilder::from(schema.fields());
         let mut columns = record_batch.columns().to_vec();
-
-        let geometry_array = builder.finish();
-        columns.push(geometry_array.to_array_ref());
-        schema_builder.push(geometry_array.data_type().to_field("geometry", true));
-
         for (key, geometry_builder) in geometry_builders {
             let geometry_array = geometry_builder.finish();
             columns.push(geometry_array.to_array_ref());
             schema_builder.push(geometry_array.data_type().to_field(key, true));
         }
-
-        // Build the table
-        let schema = Arc::new(schema_builder.finish().with_metadata(metadata));
+        let _ = schema_builder
+            .metadata_mut()
+            .insert(VERSION_KEY.to_string(), VERSION.into());
+        let schema = Arc::new(schema_builder.finish());
         let record_batch = RecordBatch::try_new(schema.clone(), columns)?;
         Ok(Table {
             record_batches: vec![record_batch],
@@ -266,14 +234,11 @@ pub fn with_native_geometry(
         );
         let geometry_array = geoarrow_array::cast::from_wkb(
             &wkb_array,
-            GeoArrowType::Geometry(GeometryType::new(
-                CoordType::Interleaved,
-                Metadata::default().into(),
-            )),
+            GeoArrowType::Geometry(GeometryType::new(Metadata::default().into())),
         )?;
         let mut columns = record_batch.columns().to_vec();
         let mut schema_builder = SchemaBuilder::from(&*record_batch.schema());
-        schema_builder.push(geometry_array.data_type().to_field("geometry", true));
+        schema_builder.push(geometry_array.data_type().to_field(column_name, true));
         let schema = schema_builder.finish();
         columns.push(geometry_array.to_array_ref());
         record_batch = RecordBatch::try_new(schema.into(), columns)?;
@@ -290,7 +255,7 @@ pub fn with_wkb_geometry(mut record_batch: RecordBatch, column_name: &str) -> Re
         )?;
         let mut columns = record_batch.columns().to_vec();
         let mut schema_builder = SchemaBuilder::from(&*record_batch.schema());
-        schema_builder.push(wkb_array.data_type().to_field("geometry", true));
+        schema_builder.push(wkb_array.data_type().to_field(column_name, true));
         let schema = schema_builder.finish();
         columns.push(wkb_array.to_array_ref());
         record_batch = RecordBatch::try_new(schema.into(), columns)?;
@@ -316,14 +281,37 @@ pub fn add_wkb_metadata(mut record_batch: RecordBatch, column_name: &str) -> Res
     Ok(record_batch)
 }
 
-// We only run tests when the geoparquet feature is enabled so that we don't
-// have to add geoarrow as a dev dependency for all builds.
-#[cfg(all(test, feature = "geoparquet"))]
+fn convert_bbox(bbox: Value) -> Result<Value> {
+    let bbox = bbox
+        .as_array()
+        .expect("STAC items should always have a list as their bbox");
+    if bbox.len() == 4 {
+        Ok(json!({
+            "xmin": bbox[0].as_number().expect("all bbox values should be a number"),
+            "ymin": bbox[1].as_number().expect("all bbox values should be a number"),
+            "xmax": bbox[2].as_number().expect("all bbox values should be a number"),
+            "ymax": bbox[3].as_number().expect("all bbox values should be a number"),
+        }))
+    } else if bbox.len() == 6 {
+        Ok(json!({
+            "xmin": bbox[0].as_number().expect("all bbox values should be a number"),
+            "ymin": bbox[1].as_number().expect("all bbox values should be a number"),
+            "zmin": bbox[2].as_number().expect("all bbox values should be a number"),
+            "xmax": bbox[3].as_number().expect("all bbox values should be a number"),
+            "ymax": bbox[4].as_number().expect("all bbox values should be a number"),
+            "zmax": bbox[5].as_number().expect("all bbox values should be a number"),
+        }))
+    } else {
+        Err(Error::InvalidBbox(
+            bbox.iter().filter_map(|v| v.as_f64()).collect(),
+        ))
+    }
+}
+
+#[cfg(test)]
 mod tests {
     use super::Table;
     use crate::{Item, ItemCollection};
-    use geoparquet::GeoParquetRecordBatchReaderBuilder;
-    use std::fs::File;
 
     #[test]
     fn to_table() {
@@ -339,17 +327,6 @@ mod tests {
         let _ = table.schema().field_with_name("type").unwrap();
     }
 
-    #[test]
-    fn from_table() {
-        let file = File::open("data/extended-item.parquet").unwrap();
-        let reader = GeoParquetRecordBatchReaderBuilder::try_new(file)
-            .unwrap()
-            .build()
-            .unwrap();
-        let item_collection = super::from_record_batch_reader(reader).unwrap();
-        assert_eq!(item_collection.items.len(), 1);
-    }
-
     #[test]
     fn roundtrip() {
         let item: Item = crate::read("examples/simple-item.json").unwrap();
diff --git a/crates/core/src/geoparquet.rs b/crates/core/src/geoparquet.rs
@@ -1,12 +1,16 @@
 //! Read data from and write data in [stac-geoparquet](https://github.com/stac-utils/stac-geoparquet/blob/main/spec/stac-geoparquet-spec.md).
 
 use crate::{
-    Catalog, Collection, Item, ItemCollection, Result, Value,
+    Catalog, Collection, Error, Item, ItemCollection, Result, Value,
     geoarrow::{Table, VERSION, VERSION_KEY},
 };
 use bytes::Bytes;
-use geoparquet::{GeoParquetRecordBatchReaderBuilder, GeoParquetWriterOptions};
+use geoparquet::{
+    reader::{GeoParquetReaderBuilder, GeoParquetRecordBatchReader},
+    writer::GeoParquetWriterOptions,
+};
 use parquet::{
+    arrow::arrow_reader::ParquetRecordBatchReaderBuilder,
     file::{properties::WriterProperties, reader::ChunkReader},
     format::KeyValue,
 };
@@ -29,7 +33,15 @@ pub fn from_reader<R>(reader: R) -> Result<ItemCollection>
 where
     R: ChunkReader + 'static,
 {
-    let reader = GeoParquetRecordBatchReaderBuilder::try_new(reader)?.build()?;
+    let builder = ParquetRecordBatchReaderBuilder::try_new(reader)?;
+    let geoparquet_metadata = builder
+        .geoparquet_metadata()
+        .transpose()?
+        .ok_or(Error::MissingGeoparquetMetadata)?;
+    let geoarrow_schema =
+        builder.geoarrow_schema(&geoparquet_metadata, true, Default::default())?;
+    let reader = builder.build()?;
+    let reader = GeoParquetRecordBatchReader::try_new(reader, geoarrow_schema)?;
     crate::geoarrow::from_record_batch_reader(reader)
 }
 
@@ -121,7 +133,7 @@ where
         options.primary_column = Some("geometry".to_string());
     }
     let table = Table::from_item_collection(item_collection)?;
-    geoparquet::write_geoparquet(Box::new(table.into_reader()), writer, &options)?;
+    geoparquet::writer::write_geoparquet(Box::new(table.into_reader()), writer, &options)?;
     Ok(())
 }
 /// Create a STAC object from geoparquet data.
