feat(vortex-geo): geometry Bbox zone-map statistic + distance-filter pruning

vortex-array/src/aggregate_fn/plugin.rs

@@ -10,6 +10,7 @@ use crate::aggregate_fn::AggregateFn;
 use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnRef;
 use crate::aggregate_fn::AggregateFnVTable;
+use crate::dtype::DType;
 
 /// Reference-counted pointer to an aggregate function plugin.
 pub type AggregateFnPluginRef = Arc<dyn AggregateFnPlugin>;
@@ -28,6 +29,19 @@ pub trait AggregateFnPlugin: 'static + Send + Sync {
     /// Deserialize an aggregate function from serialized metadata.
     fn deserialize(&self, metadata: &[u8], session: &VortexSession)
     -> VortexResult<AggregateFnRef>;
+
+    /// The default zone statistic (per-chunk) for a column of `input_dtype`, or `None` if the
+    /// dtype is not supported (or this aggregate is not a zone statistic at all).
+    ///
+    /// This is how a registered aggregate volunteers itself as a per-chunk statistic: when a
+    /// zoned layout writer opens a column, it collects every plugin's default via
+    /// [`AggregateFnSession::zone_stat_defaults`], so new statistics can be added without the
+    /// writer knowing about them.
+    ///
+    /// [`AggregateFnSession::zone_stat_defaults`]: crate::aggregate_fn::session::AggregateFnSession::zone_stat_defaults
+    fn zone_stat_default(&self, _input_dtype: &DType) -> Option<AggregateFnRef> {
+        None
+    }
 }
 
 impl std::fmt::Debug for dyn AggregateFnPlugin {
@@ -51,4 +65,8 @@ impl<V: AggregateFnVTable> AggregateFnPlugin for V {
         let options = AggregateFnVTable::deserialize(self, metadata, session)?;
         Ok(AggregateFn::new(self.clone(), options).erased())
     }
+
+    fn zone_stat_default(&self, input_dtype: &DType) -> Option<AggregateFnRef> {
+        AggregateFnVTable::zone_stat_default(self, input_dtype)
+    }
 }

vortex-array/src/aggregate_fn/session.rs

@@ -10,6 +10,7 @@ use vortex_session::SessionVar;
 
 use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnPluginRef;
+use crate::aggregate_fn::AggregateFnRef;
 use crate::aggregate_fn::AggregateFnVTable;
 use crate::aggregate_fn::fns::all_nan::AllNan;
 use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct;
@@ -44,6 +45,7 @@ use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
 use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
 use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
 use crate::arrays::dict::compute::min_max::DictMinMaxKernel;
+use crate::dtype::DType;
 
 /// Session state for aggregate functions and encoding-specific aggregate kernels.
 ///
@@ -134,6 +136,22 @@ impl AggregateFnSession {
         self.registry.insert(id, pluginref);
     }
 
+    /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every
+    /// registered aggregate's `zone_stat_default`.
+    ///
+    /// Each call scans the whole plugin registry, so this is intended to be called once per
+    /// column when a zoned writer is opened, not per chunk or per row.
+    pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec<AggregateFnRef> {
+        self.registry.read(|registry| {
+            let mut fns: Vec<AggregateFnRef> = registry
+                .values()
+                .filter_map(|plugin| plugin.zone_stat_default(input_dtype))
+                .collect();
+            fns.sort_by_key(|f| f.id());
+            fns
+        })
+    }
+
     /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
     ///
     /// Lookup first checks for a kernel registered for the exact aggregate function, then falls

vortex-array/src/aggregate_fn/vtable.rs

@@ -93,6 +93,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
     /// Returns `None` if the aggregate function cannot be applied to the input dtype.
     fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType>;
 
+    /// If this aggregate should be computed as a default zone statistic for `input_dtype`, return
+    /// the bound aggregate to store. Default: not a zone-map default.
+    fn zone_stat_default(&self, _input_dtype: &DType) -> Option<AggregateFnRef> {
+        None
+    }
+
     /// DType of the intermediate partial accumulator state.
     ///
     /// Use a struct dtype when multiple fields are needed

vortex-bench/Cargo.toml

@@ -35,6 +35,8 @@ async-trait = { workspace = true }
 bzip2 = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
 futures = { workspace = true }
+geo = { workspace = true }
+geo-traits = { workspace = true }
 geoarrow = { workspace = true }
 geoarrow-cast = { workspace = true }
 get_dir = { workspace = true }

vortex-bench/src/spatialbench/benchmark.rs

@@ -96,6 +96,16 @@ impl Benchmark for SpatialBenchBenchmark {
                 crate::conversions::add_geoparquet_metadata(&zone_file, &geo).await?;
             }
         }
+
+        // Cluster the source parquet along a spatial curve so all lanes read the same layout and
+        // the geometry zone-map prune can skip chunks.
+        let derived_dirs = [
+            base_data_dir.join(Format::OnDiskVortex.name()),
+            base_data_dir.join(Format::VortexCompact.name()),
+            base_data_dir.join(NATIVE_DIR),
+        ];
+        datagen::spatially_sort_tables(&base_data_dir.join(Format::Parquet.name()), &derived_dirs)
+            .await?;
         Ok(())
     }
 

vortex-bench/src/spatialbench/datagen/mod.rs

@@ -6,9 +6,11 @@
 //! source of truth for the base tables both stages share.
 
 pub mod native;
+pub mod spatial_sort;
 pub mod table;
 pub mod wkb;
 
 pub use native::write_native_vortex;
+pub use spatial_sort::spatially_sort_tables;
 pub use table::Table;
 pub use wkb::generate_tables;

vortex-bench/src/spatialbench/datagen/spatial_sort.rs

@@ -0,0 +1,255 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Spatial clustering of the source parquet, in place, so every downstream lane (parquet,
+//! vortex-WKB, `vortex-geo-native`) reads the same layout.
+//!
+//! The geometry zone-map prune skips a chunk only when its bounding box is disjoint from the query
+//! region. In generation order every chunk's box spans the whole map, so nothing prunes; ordering
+//! rows on a Z-order (Morton) curve of each geometry's bounding-box center gives each chunk a
+//! compact box instead. The center works uniformly for every geometry type.
+//!
+//! Every table with a geometry column is sorted by its first one, each part independently — global
+//! order is unnecessary for per-chunk pruning. A parquet marker makes it idempotent; derived vortex
+//! files from pre-sort parquet are deleted so the existence-keyed conversions regenerate.
+
+use std::path::Path;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::Context;
+use arrow_array::Array;
+use arrow_array::ArrayRef;
+use arrow_array::RecordBatch;
+use arrow_array::UInt64Array;
+use arrow_schema::DataType;
+use arrow_select::concat::concat_batches;
+use arrow_select::take::take;
+use futures::TryStreamExt;
+use geo::BoundingRect;
+use geo::Geometry;
+use geo_traits::to_geo::ToGeoGeometry;
+use geoarrow::array::GenericWkbArray;
+use geoarrow::array::GeoArrowArrayAccessor;
+use geoarrow::array::WkbViewArray;
+use geoarrow::datatypes::Crs;
+use geoarrow::datatypes::Metadata;
+use geoarrow::datatypes::WkbType;
+use parquet::arrow::AsyncArrowWriter;
+use parquet::arrow::ParquetRecordBatchStreamBuilder;
+use parquet::basic::Compression;
+use parquet::file::metadata::KeyValue;
+use parquet::file::properties::WriterProperties;
+use tokio::fs::File as TokioFile;
+use tracing::info;
+
+use super::table::Table;
+use super::wkb::geo_parquet_metadata;
+
+/// Parquet metadata marker: this file is already spatially sorted (makes the step idempotent).
+const SORTED_KEY: &str = "vortex_spatial_sorted";
+
+fn geo_metadata() -> Arc<Metadata> {
+    Arc::new(Metadata::new(Crs::default(), None))
+}
+
+/// Spatially sort every table with a geometry column, in place.
+pub async fn spatially_sort_tables(
+    parquet_dir: &Path,
+    derived_dirs: &[PathBuf],
+) -> anyhow::Result<()> {
+    for table in Table::ALL.into_iter().filter(|table| table.is_generated()) {
+        sort_table_parquet(table, parquet_dir, derived_dirs).await?;
+    }
+    Ok(())
+}
+
+/// Sort every unsorted `{table}_*.parquet` part by its first geometry column's bounding-box center.
+async fn sort_table_parquet(
+    table: Table,
+    parquet_dir: &Path,
+    derived_dirs: &[PathBuf],
+) -> anyhow::Result<()> {
+    let Some(geom) = table.geometry_columns().first() else {
+        return Ok(());
+    };
+
+    let pattern = parquet_dir.join(format!("{}_*.parquet", table.name()));
+    let mut files: Vec<PathBuf> =
+        glob::glob(&pattern.to_string_lossy())?.collect::<Result<_, _>>()?;
+    files.sort();
+    let mut pending = Vec::new();
+    for file in files {
+        if !is_sorted(&file).await? {
+            pending.push(file);
+        }
+    }
+    if pending.is_empty() {
+        return Ok(());
+    }
+
+    // Delete before rewriting: an interrupted run leaves unsorted parts unmarked and re-deletes.
+    for dir in derived_dirs {
+        let stale = dir.join(format!("{}_*.vortex", table.name()));
+        for file in glob::glob(&stale.to_string_lossy())?.flatten() {
+            tokio::fs::remove_file(&file).await?;
+            info!(path = %file.display(), "removed stale derived file from pre-sort parquet");
+        }
+    }
+
+    for file in pending {
+        sort_part(&file, table, geom.name).await?;
+    }
+    Ok(())
+}
+
+/// Whether `path` already carries the [`SORTED_KEY`] marker.
+async fn is_sorted(path: &Path) -> anyhow::Result<bool> {
+    let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?;
+    Ok(builder
+        .metadata()
+        .file_metadata()
+        .key_value_metadata()
+        .is_some_and(|kvs| kvs.iter().any(|kv| kv.key == SORTED_KEY)))
+}
+
+/// Rewrite one parquet part with its rows in Z-order of `geom_col`.
+async fn sort_part(path: &Path, table: Table, geom_col: &str) -> anyhow::Result<()> {
+    let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?;
+    let schema = Arc::clone(builder.schema());
+    let mut reader = builder.build()?;
+    let mut batches = Vec::new();
+    while let Some(batch) = reader.try_next().await? {
+        batches.push(batch);
+    }
+    if batches.iter().all(|batch| batch.num_rows() == 0) {
+        return Ok(());
+    }
+    let batch = concat_batches(&schema, &batches)?;
+    drop(batches);
+
+    let geom_idx = schema.index_of(geom_col)?;
+    let (xs, ys) = wkb_centers(batch.column(geom_idx).as_ref())
+        .with_context(|| format!("decoding geometry column {geom_col} of {}", path.display()))?;
+    let indices = morton_sort_indices(&xs, &ys);
+
+    let columns: Vec<ArrayRef> = batch
+        .columns()
+        .iter()
+        .map(|column| take(column.as_ref(), &indices, None))
+        .collect::<Result<_, _>>()?;
+    let sorted = RecordBatch::try_new(Arc::clone(&schema), columns)?;
+    let num_rows = sorted.num_rows();
+
+    let tmp_path = path.with_extension("parquet.sorttmp");
+    let props = WriterProperties::builder()
+        .set_compression(Compression::SNAPPY)
+        .build();
+    let mut writer =
+        AsyncArrowWriter::try_new(TokioFile::create(&tmp_path).await?, schema, Some(props))?;
+    writer.write(&sorted).await?;
+    // A fresh write drops metadata: re-tag geo so DuckDB reads `GEOMETRY`, and add the sorted marker.
+    if let Some(geo) = geo_parquet_metadata(table) {
+        writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo)));
+    }
+    writer.append_key_value_metadata(KeyValue::new(
+        SORTED_KEY.to_string(),
+        Some(format!("morton:{geom_col}")),
+    ));
+    writer.close().await?;
+    tokio::fs::rename(&tmp_path, path).await?;
+
+    info!(
+        path = %path.display(),
+        rows = num_rows,
+        column = geom_col,
+        "spatially sorted parquet (morton z-order)"
+    );
+    Ok(())
+}
+
+/// The bounding-box center `(x, y)` of every geometry in a WKB column, whatever its geometry type.
+fn wkb_centers(column: &dyn Array) -> anyhow::Result<(Vec<f64>, Vec<f64>)> {
+    let wkb_type = WkbType::new(geo_metadata());
+    // Expanded per concrete WKB array type.
+    macro_rules! centers {
+        ($array:expr) => {{
+            let mut xy = Vec::new();
+            for item in $array.iter() {
+                xy.push(match item {
+                    Some(geometry) => bbox_center(&geometry?.to_geometry()),
+                    None => (f64::NAN, f64::NAN),
+                });
+            }
+            Ok(xy.into_iter().unzip())
+        }};
+    }
+    match column.data_type() {
+        DataType::Binary => centers!(GenericWkbArray::<i32>::try_from((column, wkb_type))?),
+        DataType::LargeBinary => centers!(GenericWkbArray::<i64>::try_from((column, wkb_type))?),
+        DataType::BinaryView => centers!(WkbViewArray::try_from((column, wkb_type))?),
+        other => anyhow::bail!("unsupported WKB column type {other}"),
+    }
+}
+
+/// The center of a geometry's bounding box, or `NaN` for an empty geometry.
+fn bbox_center(geometry: &Geometry<f64>) -> (f64, f64) {
+    geometry
+        .bounding_rect()
+        .map(|r| ((r.min().x + r.max().x) / 2.0, (r.min().y + r.max().y) / 2.0))
+        .unwrap_or((f64::NAN, f64::NAN))
+}
+
+/// Row indices ordering `(x, y)` along a 32-bit-per-axis Z-order curve, with each axis
+/// quantized over the data's own extent.
+fn morton_sort_indices(xs: &[f64], ys: &[f64]) -> UInt64Array {
+    let (mut xmin, mut xmax) = (f64::INFINITY, f64::NEG_INFINITY);
+    let (mut ymin, mut ymax) = (f64::INFINITY, f64::NEG_INFINITY);
+    for (&x, &y) in xs.iter().zip(ys) {
+        if !x.is_nan() && !y.is_nan() {
+            xmin = xmin.min(x);
+            xmax = xmax.max(x);
+            ymin = ymin.min(y);
+            ymax = ymax.max(y);
+        }
+    }
+
+    let keys: Vec<u64> = xs
+        .iter()
+        .zip(ys)
+        .map(|(&x, &y)| {
+            if x.is_nan() || y.is_nan() {
+                0
+            } else {
+                interleave(quantize(x, xmin, xmax)) | (interleave(quantize(y, ymin, ymax)) << 1)
+            }
+        })
+        .collect();
+
+    let mut order: Vec<usize> = (0..xs.len()).collect();
+    order.sort_unstable_by_key(|&i| keys[i]);
+    UInt64Array::from_iter_values(order.into_iter().map(|i| i as u64))
+}
+
+/// Map `value` within `[min, max]` to the full `u32` range.
+// `t` is clamped to [0, 1], so the product never exceeds `u32::MAX`.
+#[expect(clippy::cast_possible_truncation)]
+fn quantize(value: f64, min: f64, max: f64) -> u32 {
+    let range = max - min;
+    if range <= 0.0 {
+        return 0;
+    }
+    let t = ((value - min) / range).clamp(0.0, 1.0);
+    (t * f64::from(u32::MAX)) as u32
+}
+
+/// Spread a 32-bit value's bits into the even positions of a `u64`.
+fn interleave(value: u32) -> u64 {
+    let mut n = u64::from(value);
+    n = (n | (n << 16)) & 0x0000_FFFF_0000_FFFF;
+    n = (n | (n << 8)) & 0x00FF_00FF_00FF_00FF;
+    n = (n | (n << 4)) & 0x0F0F_0F0F_0F0F_0F0F;
+    n = (n | (n << 2)) & 0x3333_3333_3333_3333;
+    n = (n | (n << 1)) & 0x5555_5555_5555_5555;
+    n
+}