feat(rust/sedona-geoparquet): GeoParquet 1.1 write support (#175)

Co-authored-by: Copilot <[email protected]>

Author: paleolimbot

Cargo.lock

@@ -16,7 +16,7 @@
 # under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Union, Optional, Any, Iterable
+from typing import TYPE_CHECKING, Union, Optional, Any, Iterable, Literal
 
 from sedonadb.utility import sedona  # noqa: F401
 
@@ -295,13 +295,15 @@ def to_parquet(
         partition_by: Optional[Union[str, Iterable[str]]] = None,
         sort_by: Optional[Union[str, Iterable[str]]] = None,
         single_file_output: Optional[bool] = None,
+        geoparquet_version: Literal["1.0", "1.1"] = "1.0",
+        overwrite_bbox_columns: bool = False,
     ):
         """Write this DataFrame to one or more (Geo)Parquet files
 
         For input that contains geometry columns, GeoParquet metadata is written
         such that suitable readers can recreate Geometry/Geography types when
-        reading the output.
-
+        reading the output and potentially read fewer row groups when only a
+        subset of the file is needed for a given query.
 
         Args:
             path: A filename or directory to which parquet file(s) should be written.
@@ -313,6 +315,21 @@ def to_parquet(
                 file vs. writing one file per partition to a directory. By default,
                 a single file is written if `partition_by` is unspecified and
                 `path` ends with `.parquet`.
+            geoparquet_version: GeoParquet metadata version to write if output contains
+                one or more geometry columns. The default (1.0) is the most widely
+                supported and will result in geometry columns being recognized in many
+                readers; however, only includes statistics at the file level.
+
+                Use GeoParquet 1.1 to compute an additional bounding box column
+                for every geometry column in the output: some readers can use these columns
+                to prune row groups when files contain an effective spatial ordering.
+                The extra columns will appear just before their geometry column and
+                will be named "[geom_col_name]_bbox" for all geometry columns except
+                "geometry", whose bounding box column name is just "bbox".
+            overwrite_bbox_columns: Use `True` to overwrite any bounding box columns
+                that already exist in the input. This is useful in a read -> modify
+                -> write scenario to ensure these columns are up-to-date. If `False`
+                (the default), an error will be raised if a bbox column already exists.
 
         Examples:
 
@@ -344,7 +361,13 @@ def to_parquet(
             sort_by = []
 
         self._impl.to_parquet(
-            self._ctx, str(path), partition_by, sort_by, single_file_output
+            self._ctx,
+            str(path),
+            partition_by,
+            sort_by,
+            single_file_output,
+            geoparquet_version,
+            overwrite_bbox_columns,
         )
 
     def show(

python/sedonadb/python/sedonadb/dataframe.py

@@ -32,7 +32,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
 use sedona::context::{SedonaDataFrame, SedonaWriteOptions};
 use sedona::show::{DisplayMode, DisplayTableOptions};
-use sedona_geoparquet::options::TableGeoParquetOptions;
+use sedona_geoparquet::options::{GeoParquetVersion, TableGeoParquetOptions};
 use sedona_schema::schema::SedonaSchema;
 use tokio::runtime::Runtime;
 
@@ -139,6 +139,7 @@ impl InternalDataFrame {
         ))
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn to_parquet<'py>(
         &self,
         py: Python<'py>,
@@ -147,6 +148,8 @@ impl InternalDataFrame {
         partition_by: Vec<String>,
         sort_by: Vec<String>,
         single_file_output: bool,
+        geoparquet_version: Option<String>,
+        overwrite_bbox_columns: bool,
     ) -> Result<(), PySedonaError> {
         // sort_by needs to be SortExpr. A Vec<String> can unambiguously be interpreted as
         // field names (ascending), but other types of expressions aren't supported here yet.
@@ -162,7 +165,14 @@ impl InternalDataFrame {
             .with_partition_by(partition_by)
             .with_sort_by(sort_by_expr)
             .with_single_file_output(single_file_output);
-        let writer_options = TableGeoParquetOptions::default();
+
+        let mut writer_options = TableGeoParquetOptions::new();
+        writer_options.overwrite_bbox_columns = overwrite_bbox_columns;
+        if let Some(geoparquet_version) = geoparquet_version {
+            writer_options.geoparquet_version = geoparquet_version.parse()?;
+        } else {
+            writer_options.geoparquet_version = GeoParquetVersion::Omitted;
+        }
 
         wait_for_future(
             py,

python/sedonadb/src/dataframe.rs

@@ -15,14 +15,17 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import pytest
+import json
 import tempfile
-import shapely
+from pathlib import Path
+
 import geopandas
 import geopandas.testing
+import pytest
+import shapely
 from pyarrow import parquet
-from pathlib import Path
-from sedonadb.testing import geom_or_null, SedonaDB, DuckDB, skip_if_not_exists
+from sedonadb._lib import SedonaError
+from sedonadb.testing import DuckDB, SedonaDB, geom_or_null, skip_if_not_exists
 
 
 @pytest.mark.parametrize("name", ["water-junc", "water-point"])
@@ -257,6 +260,68 @@ def test_write_geoparquet_geometry(con, geoarrow_data, name):
         geopandas.testing.assert_geodataframe_equal(gdf_roundtrip, gdf)
 
 
+def test_write_geoparquet_1_1(con, geoarrow_data):
+    # Checks GeoParquet 1.1 support specifically
+    path = geoarrow_data / "ns-water" / "files" / "ns-water_water-junc_geo.parquet"
+    skip_if_not_exists(path)
+
+    gdf = geopandas.read_parquet(path).sort_values(by="OBJECTID").reset_index(drop=True)
+
+    with tempfile.TemporaryDirectory() as td:
+        tmp_parquet = Path(td) / "tmp.parquet"
+        con.create_data_frame(gdf).to_parquet(
+            tmp_parquet, sort_by="OBJECTID", geoparquet_version="1.1"
+        )
+
+        file_kv_metadata = parquet.ParquetFile(tmp_parquet).metadata.metadata
+        assert b"geo" in file_kv_metadata
+        geo_metadata = json.loads(file_kv_metadata[b"geo"])
+        assert geo_metadata["version"] == "1.1.0"
+        geo_column = geo_metadata["columns"]["geometry"]
+        assert geo_column["covering"] == {
+            "bbox": {
+                "xmin": ["bbox", "xmin"],
+                "ymin": ["bbox", "ymin"],
+                "xmax": ["bbox", "xmax"],
+                "ymax": ["bbox", "ymax"],
+            }
+        }
+
+        # This should still roundtrip through GeoPandas because GeoPandas removes
+        # the bbox column on read
+        gdf_roundtrip = geopandas.read_parquet(tmp_parquet)
+        assert all(gdf.columns == gdf_roundtrip.columns)
+        geopandas.testing.assert_geodataframe_equal(gdf_roundtrip, gdf)
+
+        # ...but the bbox column should still be there
+        df_roundtrip = con.read_parquet(tmp_parquet).to_pandas()
+        assert "bbox" in df_roundtrip.columns
+
+        # An attempt to rewrite this should fail because it would have to overwrite
+        # the bbox column
+        tmp_parquet2 = Path(td) / "tmp2.parquet"
+        with pytest.raises(
+            SedonaError, match="Can't overwrite GeoParquet 1.1 bbox column 'bbox'"
+        ):
+            con.read_parquet(tmp_parquet).to_parquet(
+                tmp_parquet2, geoparquet_version="1.1"
+            )
+
+        # ...unless we pass the appropriate option
+        con.read_parquet(tmp_parquet).to_parquet(
+            tmp_parquet2, geoparquet_version="1.1", overwrite_bbox_columns=True
+        )
+        df_roundtrip = con.read_parquet(tmp_parquet2).to_pandas()
+        assert "bbox" in df_roundtrip.columns
+
+
+def test_write_geoparquet_unknown(con):
+    with pytest.raises(SedonaError, match="Unexpected GeoParquet version string"):
+        con.sql("SELECT 1 as one").to_parquet(
+            "unused", geoparquet_version="not supported"
+        )
+
+
 def test_write_geoparquet_geography(con, geoarrow_data):
     # Checks a read and write of geography (rounctrip, since nobody else can read/write)
     path = (

python/sedonadb/tests/io/test_parquet.py

@@ -50,12 +50,14 @@ datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-expr = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+float_next_after = { workspace = true }
 geo-traits = { workspace = true }
 futures = { workspace = true }
 object_store = { workspace = true }
 parquet = { workspace = true }
 sedona-common = { path = "../sedona-common" }
 sedona-expr = { path = "../sedona-expr" }
+sedona-functions = { path = "../sedona-functions" }
 sedona-geometry = { path = "../sedona-geometry" }
 sedona-schema = { path = "../sedona-schema" }
 serde = { workspace = true }

rust/sedona-geoparquet/Cargo.toml

@@ -48,7 +48,7 @@ use sedona_schema::extension_type::ExtensionType;
 use crate::{
     file_opener::{storage_schema_contains_geo, GeoParquetFileOpener},
     metadata::{GeoParquetColumnEncoding, GeoParquetMetadata},
-    options::{GeoParquetVersion, TableGeoParquetOptions},
+    options::TableGeoParquetOptions,
     writer::create_geoparquet_writer_physical_plan,
 };
 use datafusion::datasource::physical_plan::ParquetSource;
@@ -91,17 +91,9 @@ impl FileFormatFactory for GeoParquetFormatFactory {
     ) -> Result<Arc<dyn FileFormat>> {
         let mut options_mut = self.options.clone().unwrap_or_default();
         let mut format_options_mut = format_options.clone();
-        options_mut.geoparquet_version =
-            if let Some(version_string) = format_options_mut.remove("geoparquet_version") {
-                match version_string.as_str() {
-                    "1.0" => GeoParquetVersion::V1_0,
-                    "1.1" => GeoParquetVersion::V1_1,
-                    "2.0" => GeoParquetVersion::V2_0,
-                    _ => GeoParquetVersion::default(),
-                }
-            } else {
-                GeoParquetVersion::default()
-            };
+        if let Some(version_string) = format_options_mut.remove("geoparquet_version") {
+            options_mut.geoparquet_version = version_string.parse()?;
+        }
 
         let inner_format = self.inner.create(state, &format_options_mut)?;
         if let Some(parquet_format) = inner_format.as_any().downcast_ref::<ParquetFormat>() {

rust/sedona-geoparquet/src/format.rs

@@ -268,6 +268,21 @@ pub struct GeoParquetCovering {
     pub bbox: GeoParquetBboxCovering,
 }
 
+impl GeoParquetCovering {
+    pub fn bbox_struct_xy(struct_column_name: &str) -> Self {
+        GeoParquetCovering {
+            bbox: GeoParquetBboxCovering {
+                xmin: vec![struct_column_name.to_string(), "xmin".to_string()],
+                ymin: vec![struct_column_name.to_string(), "ymin".to_string()],
+                zmin: None,
+                xmax: vec![struct_column_name.to_string(), "xmax".to_string()],
+                ymax: vec![struct_column_name.to_string(), "ymax".to_string()],
+                zmax: None,
+            },
+        }
+    }
+}
+
 /// Top-level GeoParquet file metadata
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct GeoParquetMetadata {

rust/sedona-geoparquet/src/metadata.rs

@@ -15,7 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::str::FromStr;
+
 use datafusion::config::TableParquetOptions;
+use datafusion_common::{plan_err, DataFusionError};
 
 /// [TableParquetOptions] wrapper with GeoParquet-specific options
 #[derive(Debug, Default, Clone)]
@@ -24,13 +27,22 @@ pub struct TableGeoParquetOptions {
     pub inner: TableParquetOptions,
     /// [GeoParquetVersion] to use when writing GeoParquet files
     pub geoparquet_version: GeoParquetVersion,
+    /// When writing [GeoParquetVersion::V1_1], use `true` to overwrite existing
+    /// bounding box columns.
+    pub overwrite_bbox_columns: bool,
+}
+
+impl TableGeoParquetOptions {
+    pub fn new() -> Self {
+        Self::default()
+    }
 }
 
 impl From<TableParquetOptions> for TableGeoParquetOptions {
     fn from(value: TableParquetOptions) -> Self {
         Self {
             inner: value,
-            geoparquet_version: GeoParquetVersion::default(),
+            ..Default::default()
         }
     }
 }
@@ -73,3 +85,19 @@ impl Default for GeoParquetVersion {
         Self::V1_0
     }
 }
+
+impl FromStr for GeoParquetVersion {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "1.0" => Ok(GeoParquetVersion::V1_0),
+            "1.1" => Ok(GeoParquetVersion::V1_1),
+            "2.0" => Ok(GeoParquetVersion::V2_0),
+            "none" => Ok(GeoParquetVersion::Omitted),
+            _ => plan_err!(
+                "Unexpected GeoParquet version string (expected '1.0', '1.1', '2.0', or 'none')"
+            ),
+        }
+    }
+}