From 119fe60559f8f330164dd6eec2693d1116bd0239 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 14 Jul 2026 17:43:27 -0400 Subject: [PATCH 1/2] feat(vortex-geo): geometry bounding-box zone-map statistic + distance-filter pruning Signed-off-by: Nemo Yu --- Cargo.lock | 3 + vortex-array/src/aggregate_fn/plugin.rs | 11 + vortex-array/src/aggregate_fn/session.rs | 15 + vortex-array/src/aggregate_fn/vtable.rs | 6 + vortex-bench/Cargo.toml | 2 + vortex-bench/src/spatialbench/benchmark.rs | 10 + vortex-bench/src/spatialbench/datagen/mod.rs | 2 + .../src/spatialbench/datagen/spatial_sort.rs | 255 +++++++++ vortex-geo/Cargo.toml | 1 + vortex-geo/src/aggregate_fn/bounds.rs | 466 ++++++++++++++++ vortex-geo/src/aggregate_fn/mod.rs | 8 + vortex-geo/src/extension/mod.rs | 32 +- vortex-geo/src/lib.rs | 12 + vortex-geo/src/prune.rs | 523 ++++++++++++++++++ vortex-geo/src/test_harness.rs | 135 ++++- vortex-layout/src/layouts/zoned/writer.rs | 18 +- 16 files changed, 1484 insertions(+), 15 deletions(-) create mode 100644 vortex-bench/src/spatialbench/datagen/spatial_sort.rs create mode 100644 vortex-geo/src/aggregate_fn/bounds.rs create mode 100644 vortex-geo/src/aggregate_fn/mod.rs create mode 100644 vortex-geo/src/prune.rs diff --git a/Cargo.lock b/Cargo.lock index 74236675f9f..ca59b40e9dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9762,6 +9762,8 @@ dependencies = [ "bzip2", "clap", "futures", + "geo", + "geo-traits", "geoarrow", "geoarrow-cast", "get_dir", @@ -10263,6 +10265,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-error", + "vortex-layout", "vortex-session", "wkb", ] diff --git a/vortex-array/src/aggregate_fn/plugin.rs b/vortex-array/src/aggregate_fn/plugin.rs index b7ff8b893ac..d07008a7338 100644 --- a/vortex-array/src/aggregate_fn/plugin.rs +++ b/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; @@ -28,6 +29,12 @@ pub trait AggregateFnPlugin: 'static + Send + Sync { /// Deserialize an aggregate function from serialized metadata. fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult; + + /// The default per-chunk zone statistic to store for a column of `input_dtype`, or `None` if + /// this aggregate isn't one. + fn zone_stat_default(&self, _input_dtype: &DType) -> Option { + None + } } impl std::fmt::Debug for dyn AggregateFnPlugin { @@ -51,4 +58,8 @@ impl 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 { + AggregateFnVTable::zone_stat_default(self, input_dtype) + } } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c496f7a9bf8..0528e88b4e6 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/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,19 @@ 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`. + pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec { + self.registry.read(|registry| { + let mut fns: Vec = 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 diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 49b28dd26d7..3f0a4cf567c 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/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; + /// 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 { + None + } + /// DType of the intermediate partial accumulator state. /// /// Use a struct dtype when multiple fields are needed diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index ff604fdb48a..f228e1571d3 100644 --- a/vortex-bench/Cargo.toml +++ b/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 } diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index f2a6016a5ac..48fa84f8ec1 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/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(()) } diff --git a/vortex-bench/src/spatialbench/datagen/mod.rs b/vortex-bench/src/spatialbench/datagen/mod.rs index 7808e06cbc3..4c0d4bd2f81 100644 --- a/vortex-bench/src/spatialbench/datagen/mod.rs +++ b/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; diff --git a/vortex-bench/src/spatialbench/datagen/spatial_sort.rs b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs new file mode 100644 index 00000000000..d77aea15846 --- /dev/null +++ b/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 { + 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 = + glob::glob(&pattern.to_string_lossy())?.collect::>()?; + 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 { + 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 = batch + .columns() + .iter() + .map(|column| take(column.as_ref(), &indices, None)) + .collect::>()?; + 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, Vec)> { + 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::::try_from((column, wkb_type))?), + DataType::LargeBinary => centers!(GenericWkbArray::::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) { + 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 = 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 = (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 +} diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 36a26bfd835..9a7980bc631 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -31,6 +31,7 @@ wkb = { workspace = true } [dev-dependencies] rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-layout = { workspace = true } [lints] workspace = true diff --git a/vortex-geo/src/aggregate_fn/bounds.rs b/vortex-geo/src/aggregate_fn/bounds.rs new file mode 100644 index 00000000000..86c0d58591a --- /dev/null +++ b/vortex-geo/src/aggregate_fn/bounds.rs @@ -0,0 +1,466 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A zone-map statistic: the 2D bounding box of a native geometry column. + +use geo::Rect as GeoRect; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::GeoMetadata; +use crate::extension::Rect; +use crate::extension::box_storage_dtype; +use crate::extension::coordinate::Dimension; +use crate::extension::flatten_coordinates; +use crate::extension::is_native_geometry; + +/// Aggregates a native geometry column's 2D minimum bounding box as the native `geoarrow.box` type. +/// Stored as a zone statistic, it lets spatial filters skip chunks whose bounding box cannot hold a +/// matching row. +#[derive(Clone, Debug)] +pub struct GeometryBounds; + +/// Running union of geometry bounding boxes, or `None` until the first row. A transient +/// `geo::Rect` value - the persisted stat is the native box (see `to_scalar`). +pub struct BoundsPartial { + rect: Option>, +} + +impl BoundsPartial { + /// Grow the accumulated box to also cover `other`. + fn merge(&mut self, other: GeoRect) { + self.rect = Some(match self.rect { + Some(cur) => GeoRect::new( + ( + cur.min().x.min(other.min().x), + cur.min().y.min(other.min().y), + ), + ( + cur.max().x.max(other.max().x), + cur.max().y.max(other.max().y), + ), + ), + None => other, + }); + } +} + +/// The stat's type: the native `geoarrow.box` (2D), nullable so an empty group is a null box. +fn bounds_dtype() -> DType { + DType::Extension( + ExtDType::::try_new(GeoMetadata::default(), bounds_storage_dtype()) + .vortex_expect("2D box storage is a valid Rect") + .erased(), + ) +} + +/// The `Rect` storage `Struct` backing the zone statistic. +fn bounds_storage_dtype() -> DType { + box_storage_dtype(Dimension::Xy, Nullability::Nullable) +} + +/// The bounding box of the raw `x`/`y` slices, or `None` when empty. +fn bounds_of(xs: &[f64], ys: &[f64]) -> Option> { + if xs.is_empty() { + return None; + } + let min_max = |vals: &[f64]| { + vals.iter() + .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| { + (lo.min(v), hi.max(v)) + }) + }; + let (xmin, xmax) = min_max(xs); + let (ymin, ymax) = min_max(ys); + Some(GeoRect::new((xmin, ymin), (xmax, ymax))) +} + +/// Read a bounds stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when +/// the scalar is null (an empty group). +fn rect_from_storage(scalar: &Scalar) -> VortexResult>> { + if scalar.is_null() { + return Ok(None); + } + let storage = scalar.as_extension().to_storage_scalar(); + let fields = storage.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("bounds missing {name}"))?, + ) + }; + Ok(Some(GeoRect::new( + (read("xmin")?, read("ymin")?), + (read("xmax")?, read("ymax")?), + ))) +} + +/// Serialize a [`GeoRect`] as a native `geoarrow.box` stat scalar (inverse of [`rect_from_storage`]). +fn rect_to_storage(rect: GeoRect) -> Scalar { + let storage = Scalar::struct_( + bounds_storage_dtype(), + vec![ + Scalar::primitive(rect.min().x, Nullability::NonNullable), + Scalar::primitive(rect.min().y, Nullability::NonNullable), + Scalar::primitive(rect.max().x, Nullability::NonNullable), + Scalar::primitive(rect.max().y, Nullability::NonNullable), + ], + ); + Scalar::extension::(GeoMetadata::default(), storage) +} + +impl AggregateFnVTable for GeometryBounds { + type Options = EmptyOptions; + type Partial = BoundsPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.geo.bounds"); + *ID + } + + // Serializable so the zoned writer can persist this as a per-chunk stat. No options to encode. + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(bounds_dtype) + } + + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(|| self.bind(EmptyOptions)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + _options: &Self::Options, + _input_dtype: &DType, + ) -> VortexResult { + Ok(BoundsPartial { rect: None }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if let Some(rect) = rect_from_storage(&other)? { + partial.merge(rect); + } + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(match partial.rect { + Some(rect) => rect_to_storage(rect), + None => Scalar::null(bounds_dtype()), + }) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.rect = None; + } + + fn is_saturated(&self, _partial: &Self::Partial) -> bool { + // A bounding box can always grow, so it is never saturated. + false + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let array = match batch { + Columnar::Canonical(canonical) => canonical.clone().into_array(), + Columnar::Constant(constant) => constant.clone().into_array(), + }; + // Min/max the raw x/y buffers directly - cheap, and avoids `to_geometry`'s panic on empty + // points (which decoding each geometry would hit). + let coords = flatten_coordinates(&array, ctx)?; + let xs = coords + .unmasked_field_by_name("x")? + .clone() + .execute::(ctx)?; + let ys = coords + .unmasked_field_by_name("y")? + .clone() + .execute::(ctx)?; + if let Some(rect) = bounds_of(xs.as_slice::(), ys.as_slice::()) { + partial.merge(rect); + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + // The stored partial is already the MBR struct, so finalizing is the identity. + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +#[cfg(test)] +mod tests { + use geo::Rect as GeoRect; + use vortex_array::ArrayRef; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::session::AggregateFnSessionExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use super::BoundsPartial; + use super::GeometryBounds; + use super::bounds_dtype; + use super::rect_from_storage; + use crate::test_harness::linestring_column; + use crate::test_harness::multilinestring_column; + use crate::test_harness::multipoint_column; + use crate::test_harness::multipolygon_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + /// One column of every native geometry type over the same `(x, y)` vertex set. + fn every_native_column(vertices: &[(f64, f64)]) -> VortexResult> { + let (xs, ys): (Vec, Vec) = vertices.iter().copied().unzip(); + let flat = vertices.to_vec(); + Ok(vec![ + point_column(xs, ys)?, + linestring_column(vec![flat.clone()])?, + multipoint_column(vec![flat.clone()])?, + polygon_column(vec![vec![flat.clone()]])?, + multilinestring_column(vec![vec![flat.clone()]])?, + multipolygon_column(vec![vec![vec![flat]]])?, + ]) + } + + /// The aggregate must be serializable so the zoned writer can persist its zone-stat descriptor. + #[test] + fn serializes_for_zone_storage() -> VortexResult<()> { + let session = vortex_array::array_session(); + let metadata = GeometryBounds + .serialize(&EmptyOptions)? + .expect("GeometryBounds must be serializable to be stored as a zone statistic"); + GeometryBounds.deserialize(&metadata, &session)?; + Ok(()) + } + + /// The MBR result's corners as `(xmin, ymin, xmax, ymax)`. + fn mbr(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { + let rect = rect_from_storage(result)?.expect("non-null bounds"); + Ok((rect.min().x, rect.min().y, rect.max().x, rect.max().y)) + } + + /// The MBR of a Point column is the min/max of its coordinates, accumulated across batches. + #[test] + fn point_bounds_across_batches() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + + acc.accumulate(&point_column(vec![1.0, 3.0], vec![2.0, 4.0])?, &mut ctx)?; + acc.accumulate(&point_column(vec![-1.0], vec![5.0])?, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0)); + Ok(()) + } + + /// The MBR of a Polygon column unions every ring vertex - exercising the `List>` + /// unwrap, not just the bare Point struct. + #[test] + fn polygon_bounds_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Two rectangles: (0,0)-(2,3) and (5,5)-(7,8). The chunk MBR is their union: (0,0)-(7,8). + let polygons = polygon_column(vec![ + vec![vec![(0.0, 0.0), (2.0, 0.0), (2.0, 3.0), (0.0, 3.0)]], + vec![vec![(5.0, 5.0), (7.0, 5.0), (7.0, 8.0), (5.0, 8.0)]], + ])?; + let dtype = polygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + acc.accumulate(&polygons, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (0.0, 0.0, 7.0, 8.0)); + Ok(()) + } + + /// Every native geometry type over the same vertex set yields the same MBR - the zone stat + /// covers the whole type family. + #[test] + fn bounds_cover_every_native_geometry_type() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + for column in every_native_column(&[(1.0, 2.0), (-1.0, 5.0), (3.0, 4.0)])? { + let mut acc = + Accumulator::try_new(GeometryBounds, EmptyOptions, column.dtype().clone())?; + acc.accumulate(&column, &mut ctx)?; + assert_eq!( + mbr(&acc.finish()?)?, + (-1.0, 2.0, 3.0, 5.0), + "MBR mismatch for {}", + column.dtype() + ); + } + Ok(()) + } + + /// The MBR of a MultiPolygon column unions every vertex of every polygon's rings - exercising + /// the triple-`List` unwrap. + #[test] + fn multipolygon_bounds_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Multipolygon 0: squares (0,0)-(1,1) and (4,4)-(5,5); multipolygon 1: square (-3,7)-(-2,9). + let multipolygons = multipolygon_column(vec![ + vec![ + vec![vec![(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]], + vec![vec![(4.0, 4.0), (5.0, 4.0), (5.0, 5.0), (4.0, 5.0)]], + ], + vec![vec![vec![ + (-3.0, 7.0), + (-2.0, 7.0), + (-2.0, 9.0), + (-3.0, 9.0), + ]]], + ])?; + let dtype = multipolygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + acc.accumulate(&multipolygons, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (-3.0, 0.0, 5.0, 9.0)); + Ok(()) + } + + /// `combine_partials` unions partial boxes - the path the zoned writer takes when a zone's + /// array is chunked. + #[test] + fn combine_partials_unions_boxes() -> VortexResult<()> { + let bbox = |xmin, ymin, xmax, ymax| BoundsPartial { + rect: Some(GeoRect::new((xmin, ymin), (xmax, ymax))), + }; + let mut partial = BoundsPartial { rect: None }; + GeometryBounds.combine_partials( + &mut partial, + GeometryBounds.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, + )?; + GeometryBounds.combine_partials( + &mut partial, + GeometryBounds.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, + )?; + assert_eq!( + mbr(&GeometryBounds.to_scalar(&partial)?)?, + (0.0, -2.0, 7.0, 3.0) + ); + Ok(()) + } + + /// A null partial (an empty group's MBR) is a no-op in `combine_partials`. + #[test] + fn combine_partials_ignores_null() -> VortexResult<()> { + let mut partial = BoundsPartial { + rect: Some(GeoRect::new((0.0, 0.0), (1.0, 1.0))), + }; + GeometryBounds.combine_partials(&mut partial, Scalar::null(bounds_dtype()))?; + assert_eq!( + mbr(&GeometryBounds.to_scalar(&partial)?)?, + (0.0, 0.0, 1.0, 1.0) + ); + Ok(()) + } + + /// All-NaN coordinates: `f64::min`/`max` skip the NaNs and `geo::Rect` normalizes the result to + /// a valid (whole-plane) box, so such a chunk is always kept. Sound - NaN-coordinate rows can + /// never satisfy `distance <= r` anyway. + #[test] + fn all_nan_coordinates_kept() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let column = point_column(vec![f64::NAN, f64::NAN], vec![f64::NAN, f64::NAN])?; + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, column.dtype().clone())?; + acc.accumulate(&column, &mut ctx)?; + + let (xmin, ymin, xmax, ymax) = mbr(&acc.finish()?)?; + assert!(xmin <= xmax && ymin <= ymax); + Ok(()) + } + + /// An empty group yields a null MBR. + #[test] + fn empty_group_is_null() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + assert!(acc.finish()?.is_null()); + Ok(()) + } + + /// After `initialize`, the registry yields a default zone statistic for every native geometry + /// type (so the zoned writer stores it) but none for ordinary numeric columns. + #[test] + fn registered_as_geometry_zone_default() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + + for column in every_native_column(&[(0.0, 0.0), (1.0, 1.0)])? { + assert!( + !session + .aggregate_fns() + .zone_stat_defaults(column.dtype()) + .is_empty(), + "a geometry zone-stat default should be discovered for {}", + column.dtype() + ); + } + let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!( + session + .aggregate_fns() + .zone_stat_defaults(&i32_dtype) + .is_empty(), + "no geometry zone-stat default should apply to numeric columns" + ); + Ok(()) + } +} diff --git a/vortex-geo/src/aggregate_fn/mod.rs b/vortex-geo/src/aggregate_fn/mod.rs new file mode 100644 index 00000000000..499355beaac --- /dev/null +++ b/vortex-geo/src/aggregate_fn/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate functions over geometry columns. + +mod bounds; + +pub use bounds::*; diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index e5f551a04e8..4bd8ae3800f 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -40,11 +40,14 @@ pub use point::*; pub use polygon::*; pub use rect::*; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; @@ -57,7 +60,7 @@ use vortex_error::vortex_err; pub use wkb::*; /// Whether `dtype` is one of the native geometry extension types the geo kernels operate on. -fn is_native_geometry(dtype: &DType) -> bool { +pub(crate) fn is_native_geometry(dtype: &DType) -> bool { dtype.as_extension_opt().is_some_and(|ext| { ext.is::() || ext.is::() @@ -85,6 +88,33 @@ pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> { Ok(()) } +/// Flatten a native geometry column into a single coordinate `Struct` containing +/// every vertex of every geometry. +pub(crate) fn flatten_coordinates( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_native_geometry(array.dtype()) { + vortex_bail!( + "geo: operand is not a native geometry extension type, was {}", + array.dtype() + ); + } + let mut node = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + while matches!(node.dtype(), DType::List(..)) { + node = node + .execute::(ctx)? + .into_listview() + .elements() + .clone(); + } + node.execute::(ctx) +} + /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. pub(crate) fn geometries( array: &ArrayRef, diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 3e695dbba65..8c0ce39046e 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -3,11 +3,14 @@ use std::sync::Arc; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::stats::session::StatsSessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; +use crate::aggregate_fn::GeometryBounds; use crate::extension::LineString; use crate::extension::MultiLineString; use crate::extension::MultiPoint; @@ -16,11 +19,14 @@ use crate::extension::Point; use crate::extension::Polygon; use crate::extension::Rect; use crate::extension::WellKnownBinary; +use crate::prune::GeoDistanceBoundsPrune; use crate::scalar_fn::contains::GeoContains; use crate::scalar_fn::distance::GeoDistance; use crate::scalar_fn::intersects::GeoIntersects; +pub mod aggregate_fn; pub mod extension; +pub mod prune; pub mod scalar_fn; #[cfg(test)] mod test_harness; @@ -59,4 +65,10 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(GeoContains); session.scalar_fns().register(GeoDistance); session.scalar_fns().register(GeoIntersects); + + // The bounding-box aggregate; self-declares as a per-chunk zone stat for geometry columns. + session.aggregate_fns().register(GeometryBounds); + + // Register the spatial pruning rule that uses that bounding box. + session.stats().register_rewrite(GeoDistanceBoundsPrune); } diff --git a/vortex-geo/src/prune.rs b/vortex-geo/src/prune.rs new file mode 100644 index 00000000000..9960696c7f3 --- /dev/null +++ b/vortex-geo/src/prune.rs @@ -0,0 +1,523 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Chunk pruning for spatial filters, using the per-chunk [`GeometryBounds`] bounding box. + +use geo::BoundingRect; +use geo::Rect as GeoRect; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::expr::Expression; +use vortex_array::expr::case_when; +use vortex_array::expr::checked_add; +use vortex_array::expr::ext_storage; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::gt_eq; +use vortex_array::expr::is_root; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::lt_eq; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_array::stats::stat; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::aggregate_fn::GeometryBounds; +use crate::extension::is_native_geometry; +use crate::extension::single_geometry; +use crate::scalar_fn::distance::GeoDistance; + +/// Prunes chunks for `ST_Distance(geom, const) r` filters using the chunk's [`GeometryBounds`] +/// bounding box. Register it with `crate::initialize`; a chunk written without the `GeometryBounds` +/// statistic is scanned rather than skipped. +/// +/// All four comparisons prune: `<= r` / `< r` skip a chunk whose box is wholly beyond `r` (box +/// min-distance); `>= r` / `> r` skip one wholly within `r` (box max-distance). `==` / `!=` don't +/// prune. To add another spatial predicate, write a sibling [`StatsRewriteRule`] from the +/// `geometry_and_constant` + `distance_prune_proof` helpers; no new statistic or file-format change +/// is needed. +#[derive(Debug)] +pub struct GeoDistanceBoundsPrune; + +impl StatsRewriteRule for GeoDistanceBoundsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + // The predicate root is the comparison, not `GeoDistance`, so key on `Binary`. + Binary.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + // Only the distance comparisons prune; `==` / `!=` and other operators fall through. + let op = *expr.as_::(); + if !matches!( + op, + Operator::Lte | Operator::Lt | Operator::Gte | Operator::Gt + ) { + return Ok(None); + } + + // The left operand must be `GeoDistance(geom, const)`; the right, the radius literal. + let distance = expr.child(0); + if distance.as_opt::().is_none() { + return Ok(None); + } + let Some((geom, constant)) = geometry_and_constant(distance, ctx)? else { + return Ok(None); + }; + let Some(radius) = expr.child(1).as_opt::() else { + return Ok(None); + }; + let Ok(radius) = f64::try_from(radius) else { + return Ok(None); + }; + // A NaN radius has no sound proof: `distance NaN` is not a total order, so the chunk + // must be scanned, not pruned. + if radius.is_nan() { + return Ok(None); + } + + // Reduce `const` (any geometry type) to its bounding box. Every row sits in the chunk MBR + // and `const` in this box, so the box-to-box distance bounds the true distance soundly for + // any geometry types. + let mut exec = ctx.session().create_execution_ctx(); + let Some(query) = single_geometry(constant, &mut exec)?.bounding_rect() else { + return Ok(None); + }; + Ok(distance_prune_proof(geom, query, op, radius)) + } +} + +/// Shared bounds-pruning helper: split a symmetric geo predicate's operands into the scope-rooted +/// geometry column and the constant's scalar, or `None` when the expression doesn't have that +/// shape or the geometry's dtype has no [`GeometryBounds`] support. Symmetric only - an asymmetric +/// predicate that needs to know *which* operand is the column must recover the role separately. +fn geometry_and_constant<'a>( + expr: &'a Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + // The predicate is symmetric, so the geometry column (scope root) and the constant may be on + // either side. + let (lhs, rhs) = (expr.child(0), expr.child(1)); + let (geom, constant) = if is_root(lhs) { + (lhs, rhs) + } else if is_root(rhs) { + (rhs, lhs) + } else { + return Ok(None); + }; + + // A `GeometryBounds` stat reference only binds for dtypes it supports; anything else (e.g. a + // WKB column) must fall through to the scan. + if !is_native_geometry(&ctx.return_dtype(geom)?) { + return Ok(None); + } + + Ok(constant.as_opt::().map(|scalar| (geom, scalar))) +} + +/// Build the prune proof for `ST_Distance(geom, const) radius` from the chunk's bounding-box +/// stat, or `None` when this operator/radius cannot prune. `<=` / `<` prune a chunk whose box is +/// wholly beyond `radius` (min box-distance); `>=` / `>` prune one wholly within it (max +/// box-distance). Every row sits in the box, so those bounds prove the whole chunk one-sided. +/// +/// A distance is always `>= 0`, which decides the degenerate radii up front. +fn distance_prune_proof( + geom: &Expression, + query: GeoRect, + op: Operator, + radius: f64, +) -> Option { + // A distance is always non-negative, so degenerate radii resolve without touching the box. + match op { + // `<= r` / `< r` with a negative radius (or zero, for `<`) match nothing: prune every chunk. + Operator::Lte if radius < 0.0 => return Some(lit(true)), + Operator::Lt if radius <= 0.0 => return Some(lit(true)), + // `>= r` / `> r` with a negative radius (or zero, for `>=`) match all rows: never prune. + Operator::Gte if radius <= 0.0 => return None, + Operator::Gt if radius < 0.0 => return None, + _ => {} + } + // The stat is read through `ext_storage`/`get_item`, which propagate a missing stat (null) to + // "keep the chunk". Compared squared to avoid a `sqrt`; all operands are `>= 0`. + let mbr = ext_storage(stat(geom.clone(), GeometryBounds.bind(EmptyOptions))); + let r2 = lit(radius * radius); + Some(match op { + // Beyond the threshold: even the nearest the box can be exceeds `r`. + Operator::Lte => gt(min_dist_sq(&mbr, query), r2), + Operator::Lt => gt_eq(min_dist_sq(&mbr, query), r2), + // Within the threshold: even the farthest the box can be is below `r`. + Operator::Gte => lt(max_dist_sq(&mbr, query), r2), + Operator::Gt => lt_eq(max_dist_sq(&mbr, query), r2), + _ => return None, + }) +} + +/// Squared minimum distance between the chunk box `mbr` and the query box - a lower bound on every +/// row's distance. `dx^2 + dy^2`, each axis gap clamped at zero (zero when the intervals overlap). +fn min_dist_sq(mbr: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, mbr.clone()); + // max(0, q_lo - mbr_hi, mbr_lo - q_hi): positive only when the intervals are separated. + let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + maximum( + lit(0.0), + maximum( + binop(Operator::Sub, lit(q_lo), hi), + binop(Operator::Sub, lo, lit(q_hi)), + ), + ) + }; + let dx = gap(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = gap(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// Squared maximum distance between the chunk box `mbr` and the query box - an upper bound on every +/// row's distance. `Dx^2 + Dy^2`, each axis span the full extent of the two intervals' union. +fn max_dist_sq(mbr: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, mbr.clone()); + // max(q_hi, mbr_hi) - min(q_lo, mbr_lo): the farthest two points of the boxes can be on an axis. + // The (nullable) MBR field is passed as the second arg so `case_when`'s else branch carries the + // nullability - a missing stat then propagates null through to "keep the chunk". + let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + binop( + Operator::Sub, + maximum(lit(q_hi), hi), + minimum(lit(q_lo), lo), + ) + }; + let dx = span(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = span(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// `a b` as a binary-operator expression. +fn binop(op: Operator, a: Expression, b: Expression) -> Expression { + Binary + .try_new_expr(op, [a, b]) + .vortex_expect("binary expression") +} + +/// `e * e` as an expression. +fn square(e: Expression) -> Expression { + binop(Operator::Mul, e.clone(), e) +} + +/// `max(a, b)` as an expression. +fn maximum(a: Expression, b: Expression) -> Expression { + case_when(gt(a.clone(), b.clone()), a, b) +} + +/// `min(a, b)` as an expression. +fn minimum(a: Expression, b: Expression) -> Expression { + case_when(lt(a.clone(), b.clone()), a, b) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::expr::Expression; + use vortex_array::expr::gt_eq; + use vortex_array::expr::lit; + use vortex_array::expr::lt_eq; + use vortex_array::expr::root; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::fns::binary::Binary; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::stats::rewrite::StatsRewriteCtx; + use vortex_array::stats::rewrite::StatsRewriteRule; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_layout::layouts::zoned::zone_map::ZoneMap; + + use super::GeoDistanceBoundsPrune; + use crate::aggregate_fn::GeometryBounds; + use crate::extension::GeoMetadata; + use crate::extension::Rect; + use crate::scalar_fn::distance::GeoDistance; + use crate::test_harness::point_column; + + /// Run the rule against `GeoDistance(root, origin) radius`, operands swapped when + /// `geom_first` is false. + fn falsify_distance( + operator: Operator, + geom_first: bool, + radius: f64, + ) -> VortexResult> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let operands = if geom_first { + [root(), lit(origin)] + } else { + [lit(origin), root()] + }; + let distance = GeoDistance.new_expr(EmptyOptions, operands); + let predicate = Binary.new_expr(operator, [distance, lit(radius)]); + + GeoDistanceBoundsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + } + + /// All four distance comparisons prune (`<=`/`<` via min-distance, `>=`/`>` via max-distance); + /// `==`/`!=` are left to the scan. + #[rstest] + #[case(Operator::Lte, true)] + #[case(Operator::Lt, true)] + #[case(Operator::Gt, true)] + #[case(Operator::Gte, true)] + #[case(Operator::Eq, false)] + #[case(Operator::NotEq, false)] + fn prunes_distance_comparisons( + #[case] operator: Operator, + #[case] prunes: bool, + ) -> VortexResult<()> { + assert_eq!(falsify_distance(operator, true, 0.5)?.is_some(), prunes); + Ok(()) + } + + /// Distance is symmetric: `GeoDistance(const, geom) <= r` falsifies just like the geom-first form. + #[test] + fn falsifies_with_constant_as_left_operand() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, false, 0.5)?.is_some()); + Ok(()) + } + + /// A NaN radius must not prune - the scan's total-order compare treats `dist <= NaN` as true. + #[test] + fn nan_radius_never_prunes() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, f64::NAN)?.is_none()); + Ok(()) + } + + /// A negative radius prunes every zone - vacuously sound: `distance <= r < 0` matches no row. + #[test] + fn negative_radius_prunes_vacuously() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, -0.5)?.is_some()); + Ok(()) + } + + /// A scope dtype without `GeometryBounds` support gets no proof - the stat reference would + /// fail to bind at prune time. + #[test] + fn unsupported_scope_is_not_pruned() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let scope = DType::Primitive(PType::F64, Nullability::NonNullable); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// A comparison that does not wrap `GeoDistance` is left untouched. + #[test] + fn ignores_non_distance_comparison() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + + let predicate = lt_eq(lit(1.0f64), lit(2.0f64)); + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// End-to-end over a hand-built zone map: the far chunk is skipped, the near one kept. + #[test] + fn prunes_far_chunk_keeps_near() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + + // Two chunks' MBRs, stored as the native `geoarrow.box` stat column with default + // (unreferenced) metadata to match the aggregate's return dtype: chunk 0 near the origin + // (0,0..1,1), chunk 1 far away (100,100..101,101). + let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); + let boxes = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![ + ord(0.0, 100.0), + ord(0.0, 100.0), + ord(1.0, 101.0), + ord(1.0, 101.0), + ], + 2, + Validity::AllValid, + )? + .into_array(); + let box_dtype = + ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); + let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let zone_map = + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 2)?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + // `true` means the zone is pruned: chunk 0 (near origin) is kept, chunk 1 (far) is skipped. + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, true]); + Ok(()) + } + + /// The true-distance prune skips a chunk that is *diagonally* farther than `r`, even though + /// neither axis alone exceeds `r` - the case a per-axis box-overlap test would wrongly keep. + #[test] + fn prunes_diagonally_distant_chunk() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + + // One chunk, MBR (0.8,0.8)..(0.9,0.9): each axis is only 0.8 from the origin (<= r = 1), but + // the near corner is sqrt(0.8^2 + 0.8^2) ~= 1.13 away (> 1), so no point in the box is within 1. + let ord = |a: f64| PrimitiveArray::from_iter([a]).into_array(); + let boxes = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![ord(0.8), ord(0.8), ord(0.9), ord(0.9)], + 1, + Validity::AllValid, + )? + .into_array(); + let box_dtype = + ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); + let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let zone_map = + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 1)?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(1.0f64)); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + assert_eq!( + zone_map + .prune(&proof, &session)? + .iter() + .collect::>(), + vec![true], + ); + Ok(()) + } + + /// Backward compat: a zone map written without the `GeometryBounds` stat (an older file) keeps + /// every zone - the missing stat binds to null and `null_as_false` retains the zone. + #[test] + fn missing_bounds_stat_keeps_all_zones() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let zone_map = ZoneMap::try_new( + point_dtype.clone(), + StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?, + Arc::new([]), + 1, + 2, + )?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let proof = lt_eq(distance, lit(0.5f64)) + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false]); + Ok(()) + } + + /// A `>= r` filter prunes a chunk lying wholly *within* `r` (every row nearer than `r`, so none + /// satisfy `>= r`) via the box max-distance, while a chunk beyond `r` is kept. + #[test] + fn prunes_within_chunk_for_far_filter() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + + // Chunk 0 (MBR 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; + // chunk 1 (100,100..101,101) is entirely beyond it. + let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); + let boxes = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![ + ord(0.0, 100.0), + ord(0.0, 100.0), + ord(0.5, 101.0), + ord(0.5, 101.0), + ], + 2, + Validity::AllValid, + )? + .into_array(); + let box_dtype = + ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); + let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let zone_map = + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 2)?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let proof = gt_eq(distance, lit(2.0f64)) + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + // Chunk 0 (within 2) is pruned for `>= 2`; chunk 1 (beyond 2) is kept. + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![true, false]); + Ok(()) + } +} diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 4711e2e64db..34f9862f4b9 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -6,17 +6,35 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::extension::GeoMetadata; +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; +use crate::extension::MultiPolygon; use crate::extension::Point; +use crate::extension::Polygon; use crate::extension::Rect; +use crate::extension::box_storage_dtype; use crate::extension::coordinate::Coordinate; +use crate::extension::coordinate::Dimension; use crate::extension::coordinate::coordinate_from_struct; +use crate::extension::linestring_storage_dtype; +use crate::extension::multilinestring_storage_dtype; +use crate::extension::multipoint_storage_dtype; +use crate::extension::multipolygon_storage_dtype; +use crate::extension::polygon_storage_dtype; /// The WGS 84 (`EPSG:4326`) metadata tagged onto test geometry columns. fn wgs84() -> GeoMetadata { @@ -25,18 +43,115 @@ fn wgs84() -> GeoMetadata { } } -/// A `Point` column (CRS `EPSG:4326`) over the given x/y coordinates. -pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { - let storage = StructArray::from_fields(&[ +/// A coordinate `Struct` over the parallel x/y buffers. +fn xy_struct(xs: Vec, ys: Vec) -> VortexResult { + Ok(StructArray::from_fields(&[ ("x", PrimitiveArray::from_iter(xs).into_array()), ("y", PrimitiveArray::from_iter(ys).into_array()), ])? - .into_array(); - let dtype = ExtDType::::try_new(wgs84(), storage.dtype().clone())?; - Ok(ExtensionArray::new(dtype.erased(), storage).into_array()) + .into_array()) +} + +/// A list offset as `i32`. +fn offset(n: usize) -> VortexResult { + i32::try_from(n).map_err(|_| vortex_err!("geometry offset overflow")) +} + +/// Wrap `elements` — built from `rows` flattened one level — in a non-nullable `List` with one +/// entry per row. +fn nest(rows: &[Vec], elements: ArrayRef) -> VortexResult { + let mut offsets = vec![0i32]; + let mut len = 0usize; + for row in rows { + len += row.len(); + offsets.push(offset(len)?); + } + Ok(ListArray::try_new( + elements, + PrimitiveArray::from_iter(offsets).into_array(), + Validity::NonNullable, + )? + .into_array()) +} + +/// `List>` storage: one list of `(x, y)` vertices per row. +fn vertex_lists(rows: &[Vec<(f64, f64)>]) -> VortexResult { + let (xs, ys) = rows.iter().flatten().copied().unzip(); + nest(rows, xy_struct(xs, ys)?) +} + +/// `List>>` storage: one list of vertex lists per row. +fn vertex_list_lists(rows: &[Vec>]) -> VortexResult { + let inner: Vec> = rows.iter().flatten().cloned().collect(); + nest(rows, vertex_lists(&inner)?) +} + +/// Wrap `storage` in the geometry extension `vtable` (CRS `EPSG:4326`) with the canonical +/// `storage_dtype` of that type. +fn geo_column + Default>( + storage: ArrayRef, + storage_dtype: DType, +) -> VortexResult { + let dtype = ExtDType::::try_new(wgs84(), storage_dtype)?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) +} + +/// A `Point` column over the given x/y coordinates, stored as `Struct`. +pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { + let storage = xy_struct(xs, ys)?; + let storage_dtype = storage.dtype().clone(); + geo_column::(storage, storage_dtype) +} + +/// A `LineString` column: each line a list of `(x, y)` vertices, stored as `List>`. +pub(crate) fn linestring_column(lines: Vec>) -> VortexResult { + geo_column::( + vertex_lists(&lines)?, + linestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `MultiPoint` column: each row a list of `(x, y)` points, stored as `List>`. +pub(crate) fn multipoint_column(points: Vec>) -> VortexResult { + geo_column::( + vertex_lists(&points)?, + multipoint_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `Polygon` column: each polygon a list of rings, each ring a list of `(x, y)` vertices, +/// stored as `List>>`. +pub(crate) fn polygon_column(polygons: Vec>>) -> VortexResult { + geo_column::( + vertex_list_lists(&polygons)?, + polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `MultiLineString` column: each row a list of lines, stored as `List>>`. +pub(crate) fn multilinestring_column( + multilines: Vec>>, +) -> VortexResult { + geo_column::( + vertex_list_lists(&multilines)?, + multilinestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// One multipolygon: polygons → rings → `(x, y)` vertices. +pub(crate) type MultiPolygonRings = Vec>>; + +/// A `MultiPolygon` column, stored as `List>>>`. +pub(crate) fn multipolygon_column(multipolygons: Vec) -> VortexResult { + let polygons: Vec>> = multipolygons.iter().flatten().cloned().collect(); + geo_column::( + nest(&multipolygons, vertex_list_lists(&polygons)?)?, + multipolygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) } -/// A 2D `Rect` (`geoarrow.box`) column (CRS `EPSG:4326`) over `(xmin, ymin, xmax, ymax)` boxes. +/// A 2D `Rect` (`geoarrow.box`) column over `(xmin, ymin, xmax, ymax)` boxes, stored as +/// `Struct`. pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult { let field = |select: fn(&(f64, f64, f64, f64)) -> f64| { PrimitiveArray::from_iter(boxes.iter().map(select)).into_array() @@ -48,8 +163,10 @@ pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult::try_new(wgs84(), storage.dtype().clone())?; - Ok(ExtensionArray::new(dtype.erased(), storage).into_array()) + geo_column::( + storage, + box_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) } /// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index c1c12d7ddea..048905dade6 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -26,6 +26,7 @@ use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -109,7 +110,7 @@ impl LayoutStrategy for ZonedStrategy { .options .aggregate_fns .clone() - .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype())); + .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype(), session)); let compute_session = session.clone(); let stats_accumulator = Arc::new(Mutex::new(AggregateStatsAccumulator::new( @@ -195,7 +196,7 @@ impl LayoutStrategy for ZonedStrategy { } } -fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { +fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { let (max, min) = match dtype { DType::Utf8(_) | DType::Binary(_) => ( BoundedMax.bind(BoundedMaxOptions { @@ -221,6 +222,9 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); + // Stats from geo extension types are discovered from the registry at runtime instead. + aggregate_fns.extend(session.aggregate_fns().zone_stat_defaults(dtype)); + aggregate_fns.into() } @@ -240,7 +244,10 @@ mod tests { #[test] fn default_aggregates_bound_variable_length_min_max() { - let aggregate_fns = default_zoned_aggregate_fns(&DType::Utf8(Nullability::NonNullable)); + let aggregate_fns = default_zoned_aggregate_fns( + &DType::Utf8(Nullability::NonNullable), + &vortex_array::array_session(), + ); assert_eq!( aggregate_fns[0].as_::().max_bytes, @@ -254,7 +261,8 @@ mod tests { #[test] fn default_aggregates_keep_fixed_width_min_max_exact() { - let aggregate_fns = default_zoned_aggregate_fns(&PType::I32.into()); + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); @@ -266,7 +274,7 @@ mod tests { let dtype = DType::Extension( Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), ); - let aggregate_fns = default_zoned_aggregate_fns(&dtype); + let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); assert!( aggregate_fns From 6b67b8ebf88a9f061b504cae1b873318eb12b656 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 15 Jul 2026 11:37:31 -0400 Subject: [PATCH 2/2] refactor(vortex-geo): address review - GeometryAabb rename + AABB terminology sweep - Rename GeometryBounds -> GeometryAabb (id vortex.geo.bounds -> vortex.geo.aabb), BoundsPartial -> AabbPartial, aggregate_fn/bounds.rs -> aabb.rs: the statistic is specifically the axis-aligned bounding box, leaving room for other bbox kinds later - Rename GeoDistanceBoundsPrune -> GeoDistancePrune; the AABB is an implementation detail carried by the doc comment - Sweep docs, comments, and test names to axis-aligned bounding box (AABB) terminology - Document what zone_stat_default is for and note zone_stat_defaults' linear registry scan (intended once per column at writer open) - Present GeometryAabb as a plain aggregate whose zone-stat role is an application - Collapse AabbPartial::merge to map_or - Comment which other operators could prune in the future (==) or cannot (!=), and why the radius try_from fall-through is a sound decline rather than an error - Extract test_harness::geo_session() to replace repeated session+initialize in tests Signed-off-by: Nemo Yu --- vortex-array/src/aggregate_fn/plugin.rs | 11 +- vortex-array/src/aggregate_fn/session.rs | 3 + .../src/aggregate_fn/{bounds.rs => aabb.rs} | 148 +++++++++--------- vortex-geo/src/aggregate_fn/mod.rs | 4 +- vortex-geo/src/extension/rect.rs | 2 +- vortex-geo/src/lib.rs | 13 +- vortex-geo/src/prune.rs | 125 +++++++-------- vortex-geo/src/test_harness.rs | 8 + vortex-geo/src/tests/mod.rs | 6 +- 9 files changed, 167 insertions(+), 153 deletions(-) rename vortex-geo/src/aggregate_fn/{bounds.rs => aabb.rs} (74%) diff --git a/vortex-array/src/aggregate_fn/plugin.rs b/vortex-array/src/aggregate_fn/plugin.rs index d07008a7338..baf6f303466 100644 --- a/vortex-array/src/aggregate_fn/plugin.rs +++ b/vortex-array/src/aggregate_fn/plugin.rs @@ -30,8 +30,15 @@ pub trait AggregateFnPlugin: 'static + Send + Sync { fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult; - /// The default per-chunk zone statistic to store for a column of `input_dtype`, or `None` if - /// this aggregate isn't one. + /// 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 { None } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 0528e88b4e6..72ef78a400c 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -138,6 +138,9 @@ impl AggregateFnSession { /// 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 { self.registry.read(|registry| { let mut fns: Vec = registry diff --git a/vortex-geo/src/aggregate_fn/bounds.rs b/vortex-geo/src/aggregate_fn/aabb.rs similarity index 74% rename from vortex-geo/src/aggregate_fn/bounds.rs rename to vortex-geo/src/aggregate_fn/aabb.rs index 86c0d58591a..f217494fc1a 100644 --- a/vortex-geo/src/aggregate_fn/bounds.rs +++ b/vortex-geo/src/aggregate_fn/aabb.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A zone-map statistic: the 2D bounding box of a native geometry column. +//! The 2D axis-aligned bounding-box (AABB) aggregate for native geometry columns. use geo::Rect as GeoRect; use vortex_array::ArrayRef; @@ -32,23 +32,23 @@ use crate::extension::coordinate::Dimension; use crate::extension::flatten_coordinates; use crate::extension::is_native_geometry; -/// Aggregates a native geometry column's 2D minimum bounding box as the native `geoarrow.box` type. -/// Stored as a zone statistic, it lets spatial filters skip chunks whose bounding box cannot hold a -/// matching row. +/// Aggregates a native geometry column's 2D axis-aligned bounding box (AABB) as a native +/// `geoarrow.box`, the spatial analogue of min/max. Also the default zone statistic for such +/// columns (via `zone_stat_default`). #[derive(Clone, Debug)] -pub struct GeometryBounds; +pub struct GeometryAabb; -/// Running union of geometry bounding boxes, or `None` until the first row. A transient +/// Running union of geometry AABBs, or `None` until the first row. A transient /// `geo::Rect` value - the persisted stat is the native box (see `to_scalar`). -pub struct BoundsPartial { +pub struct AabbPartial { rect: Option>, } -impl BoundsPartial { +impl AabbPartial { /// Grow the accumulated box to also cover `other`. fn merge(&mut self, other: GeoRect) { - self.rect = Some(match self.rect { - Some(cur) => GeoRect::new( + self.rect = Some(self.rect.map_or(other, |cur| { + GeoRect::new( ( cur.min().x.min(other.min().x), cur.min().y.min(other.min().y), @@ -57,28 +57,27 @@ impl BoundsPartial { cur.max().x.max(other.max().x), cur.max().y.max(other.max().y), ), - ), - None => other, - }); + ) + })); } } /// The stat's type: the native `geoarrow.box` (2D), nullable so an empty group is a null box. -fn bounds_dtype() -> DType { +fn aabb_dtype() -> DType { DType::Extension( - ExtDType::::try_new(GeoMetadata::default(), bounds_storage_dtype()) + ExtDType::::try_new(GeoMetadata::default(), aabb_storage_dtype()) .vortex_expect("2D box storage is a valid Rect") .erased(), ) } /// The `Rect` storage `Struct` backing the zone statistic. -fn bounds_storage_dtype() -> DType { +fn aabb_storage_dtype() -> DType { box_storage_dtype(Dimension::Xy, Nullability::Nullable) } -/// The bounding box of the raw `x`/`y` slices, or `None` when empty. -fn bounds_of(xs: &[f64], ys: &[f64]) -> Option> { +/// The AABB of the raw `x`/`y` slices, or `None` when empty. +fn aabb_of(xs: &[f64], ys: &[f64]) -> Option> { if xs.is_empty() { return None; } @@ -93,7 +92,7 @@ fn bounds_of(xs: &[f64], ys: &[f64]) -> Option> { Some(GeoRect::new((xmin, ymin), (xmax, ymax))) } -/// Read a bounds stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when +/// Read an AABB stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when /// the scalar is null (an empty group). fn rect_from_storage(scalar: &Scalar) -> VortexResult>> { if scalar.is_null() { @@ -105,7 +104,7 @@ fn rect_from_storage(scalar: &Scalar) -> VortexResult>> { f64::try_from( &fields .field(name) - .ok_or_else(|| vortex_err!("bounds missing {name}"))?, + .ok_or_else(|| vortex_err!("AABB missing {name}"))?, ) }; Ok(Some(GeoRect::new( @@ -117,7 +116,7 @@ fn rect_from_storage(scalar: &Scalar) -> VortexResult>> { /// Serialize a [`GeoRect`] as a native `geoarrow.box` stat scalar (inverse of [`rect_from_storage`]). fn rect_to_storage(rect: GeoRect) -> Scalar { let storage = Scalar::struct_( - bounds_storage_dtype(), + aabb_storage_dtype(), vec![ Scalar::primitive(rect.min().x, Nullability::NonNullable), Scalar::primitive(rect.min().y, Nullability::NonNullable), @@ -128,12 +127,12 @@ fn rect_to_storage(rect: GeoRect) -> Scalar { Scalar::extension::(GeoMetadata::default(), storage) } -impl AggregateFnVTable for GeometryBounds { +impl AggregateFnVTable for GeometryAabb { type Options = EmptyOptions; - type Partial = BoundsPartial; + type Partial = AabbPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.geo.bounds"); + static ID: CachedId = CachedId::new("vortex.geo.aabb"); *ID } @@ -151,7 +150,7 @@ impl AggregateFnVTable for GeometryBounds { } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - is_native_geometry(input_dtype).then(bounds_dtype) + is_native_geometry(input_dtype).then(aabb_dtype) } fn zone_stat_default(&self, input_dtype: &DType) -> Option { @@ -167,7 +166,7 @@ impl AggregateFnVTable for GeometryBounds { _options: &Self::Options, _input_dtype: &DType, ) -> VortexResult { - Ok(BoundsPartial { rect: None }) + Ok(AabbPartial { rect: None }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { @@ -180,7 +179,7 @@ impl AggregateFnVTable for GeometryBounds { fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(match partial.rect { Some(rect) => rect_to_storage(rect), - None => Scalar::null(bounds_dtype()), + None => Scalar::null(aabb_dtype()), }) } @@ -189,7 +188,7 @@ impl AggregateFnVTable for GeometryBounds { } fn is_saturated(&self, _partial: &Self::Partial) -> bool { - // A bounding box can always grow, so it is never saturated. + // An AABB can always grow, so it is never saturated. false } @@ -214,14 +213,14 @@ impl AggregateFnVTable for GeometryBounds { .unmasked_field_by_name("y")? .clone() .execute::(ctx)?; - if let Some(rect) = bounds_of(xs.as_slice::(), ys.as_slice::()) { + if let Some(rect) = aabb_of(xs.as_slice::(), ys.as_slice::()) { partial.merge(rect); } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - // The stored partial is already the MBR struct, so finalizing is the identity. + // The stored partial is already the AABB struct, so finalizing is the identity. Ok(partials) } @@ -246,10 +245,11 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_error::VortexResult; - use super::BoundsPartial; - use super::GeometryBounds; - use super::bounds_dtype; + use super::AabbPartial; + use super::GeometryAabb; + use super::aabb_dtype; use super::rect_from_storage; + use crate::test_harness::geo_session; use crate::test_harness::linestring_column; use crate::test_harness::multilinestring_column; use crate::test_harness::multipoint_column; @@ -275,80 +275,79 @@ mod tests { #[test] fn serializes_for_zone_storage() -> VortexResult<()> { let session = vortex_array::array_session(); - let metadata = GeometryBounds + let metadata = GeometryAabb .serialize(&EmptyOptions)? - .expect("GeometryBounds must be serializable to be stored as a zone statistic"); - GeometryBounds.deserialize(&metadata, &session)?; + .expect("GeometryAabb must be serializable to be stored as a zone statistic"); + GeometryAabb.deserialize(&metadata, &session)?; Ok(()) } - /// The MBR result's corners as `(xmin, ymin, xmax, ymax)`. - fn mbr(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { - let rect = rect_from_storage(result)?.expect("non-null bounds"); + /// The AABB result's corners as `(xmin, ymin, xmax, ymax)`. + fn aabb(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { + let rect = rect_from_storage(result)?.expect("non-null AABB"); Ok((rect.min().x, rect.min().y, rect.max().x, rect.max().y)) } - /// The MBR of a Point column is the min/max of its coordinates, accumulated across batches. + /// The AABB of a Point column is the min/max of its coordinates, accumulated across batches. #[test] - fn point_bounds_across_batches() -> VortexResult<()> { + fn point_aabb_across_batches() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; acc.accumulate(&point_column(vec![1.0, 3.0], vec![2.0, 4.0])?, &mut ctx)?; acc.accumulate(&point_column(vec![-1.0], vec![5.0])?, &mut ctx)?; - assert_eq!(mbr(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0)); + assert_eq!(aabb(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0)); Ok(()) } - /// The MBR of a Polygon column unions every ring vertex - exercising the `List>` + /// The AABB of a Polygon column unions every ring vertex - exercising the `List>` /// unwrap, not just the bare Point struct. #[test] - fn polygon_bounds_union_all_vertices() -> VortexResult<()> { + fn polygon_aabb_union_all_vertices() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); - // Two rectangles: (0,0)-(2,3) and (5,5)-(7,8). The chunk MBR is their union: (0,0)-(7,8). + // Two rectangles: (0,0)-(2,3) and (5,5)-(7,8). The chunk AABB is their union: (0,0)-(7,8). let polygons = polygon_column(vec![ vec![vec![(0.0, 0.0), (2.0, 0.0), (2.0, 3.0), (0.0, 3.0)]], vec![vec![(5.0, 5.0), (7.0, 5.0), (7.0, 8.0), (5.0, 8.0)]], ])?; let dtype = polygons.dtype().clone(); - let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; acc.accumulate(&polygons, &mut ctx)?; - assert_eq!(mbr(&acc.finish()?)?, (0.0, 0.0, 7.0, 8.0)); + assert_eq!(aabb(&acc.finish()?)?, (0.0, 0.0, 7.0, 8.0)); Ok(()) } - /// Every native geometry type over the same vertex set yields the same MBR - the zone stat + /// Every native geometry type over the same vertex set yields the same AABB - the zone stat /// covers the whole type family. #[test] - fn bounds_cover_every_native_geometry_type() -> VortexResult<()> { + fn aabb_covers_every_native_geometry_type() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); for column in every_native_column(&[(1.0, 2.0), (-1.0, 5.0), (3.0, 4.0)])? { - let mut acc = - Accumulator::try_new(GeometryBounds, EmptyOptions, column.dtype().clone())?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, column.dtype().clone())?; acc.accumulate(&column, &mut ctx)?; assert_eq!( - mbr(&acc.finish()?)?, + aabb(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0), - "MBR mismatch for {}", + "AABB mismatch for {}", column.dtype() ); } Ok(()) } - /// The MBR of a MultiPolygon column unions every vertex of every polygon's rings - exercising + /// The AABB of a MultiPolygon column unions every vertex of every polygon's rings - exercising /// the triple-`List` unwrap. #[test] - fn multipolygon_bounds_union_all_vertices() -> VortexResult<()> { + fn multipolygon_aabb_union_all_vertices() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); @@ -366,10 +365,10 @@ mod tests { ]]], ])?; let dtype = multipolygons.dtype().clone(); - let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; acc.accumulate(&multipolygons, &mut ctx)?; - assert_eq!(mbr(&acc.finish()?)?, (-3.0, 0.0, 5.0, 9.0)); + assert_eq!(aabb(&acc.finish()?)?, (-3.0, 0.0, 5.0, 9.0)); Ok(()) } @@ -377,34 +376,34 @@ mod tests { /// array is chunked. #[test] fn combine_partials_unions_boxes() -> VortexResult<()> { - let bbox = |xmin, ymin, xmax, ymax| BoundsPartial { + let bbox = |xmin, ymin, xmax, ymax| AabbPartial { rect: Some(GeoRect::new((xmin, ymin), (xmax, ymax))), }; - let mut partial = BoundsPartial { rect: None }; - GeometryBounds.combine_partials( + let mut partial = AabbPartial { rect: None }; + GeometryAabb.combine_partials( &mut partial, - GeometryBounds.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, + GeometryAabb.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, )?; - GeometryBounds.combine_partials( + GeometryAabb.combine_partials( &mut partial, - GeometryBounds.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, + GeometryAabb.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, )?; assert_eq!( - mbr(&GeometryBounds.to_scalar(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&partial)?)?, (0.0, -2.0, 7.0, 3.0) ); Ok(()) } - /// A null partial (an empty group's MBR) is a no-op in `combine_partials`. + /// A null partial (an empty group's AABB) is a no-op in `combine_partials`. #[test] fn combine_partials_ignores_null() -> VortexResult<()> { - let mut partial = BoundsPartial { + let mut partial = AabbPartial { rect: Some(GeoRect::new((0.0, 0.0), (1.0, 1.0))), }; - GeometryBounds.combine_partials(&mut partial, Scalar::null(bounds_dtype()))?; + GeometryAabb.combine_partials(&mut partial, Scalar::null(aabb_dtype()))?; assert_eq!( - mbr(&GeometryBounds.to_scalar(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&partial)?)?, (0.0, 0.0, 1.0, 1.0) ); Ok(()) @@ -419,19 +418,19 @@ mod tests { let mut ctx = session.create_execution_ctx(); let column = point_column(vec![f64::NAN, f64::NAN], vec![f64::NAN, f64::NAN])?; - let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, column.dtype().clone())?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, column.dtype().clone())?; acc.accumulate(&column, &mut ctx)?; - let (xmin, ymin, xmax, ymax) = mbr(&acc.finish()?)?; + let (xmin, ymin, xmax, ymax) = aabb(&acc.finish()?)?; assert!(xmin <= xmax && ymin <= ymax); Ok(()) } - /// An empty group yields a null MBR. + /// An empty group yields a null AABB. #[test] fn empty_group_is_null() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; assert!(acc.finish()?.is_null()); Ok(()) } @@ -440,8 +439,7 @@ mod tests { /// type (so the zoned writer stores it) but none for ordinary numeric columns. #[test] fn registered_as_geometry_zone_default() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); for column in every_native_column(&[(0.0, 0.0), (1.0, 1.0)])? { assert!( diff --git a/vortex-geo/src/aggregate_fn/mod.rs b/vortex-geo/src/aggregate_fn/mod.rs index 499355beaac..70831f02905 100644 --- a/vortex-geo/src/aggregate_fn/mod.rs +++ b/vortex-geo/src/aggregate_fn/mod.rs @@ -3,6 +3,6 @@ //! Aggregate functions over geometry columns. -mod bounds; +mod aabb; -pub use bounds::*; +pub use aabb::*; diff --git a/vortex-geo/src/extension/rect.rs b/vortex-geo/src/extension/rect.rs index d2dddf51e79..3576966368c 100644 --- a/vortex-geo/src/extension/rect.rs +++ b/vortex-geo/src/extension/rect.rs @@ -56,7 +56,7 @@ use super::coordinate::Dimension; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; -/// A bounding box (`geoarrow.box`), stored as `Struct`. +/// An axis-aligned bounding box (`geoarrow.box`), stored as `Struct`. // Named `Rect`, not `Box`: matches `geo::Rect` / geoarrow-rs `RectArray`, and `Box` is a std name. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct Rect; diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 8c0ce39046e..652311be3b2 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -10,7 +10,7 @@ use vortex_array::stats::session::StatsSessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; -use crate::aggregate_fn::GeometryBounds; +use crate::aggregate_fn::GeometryAabb; use crate::extension::LineString; use crate::extension::MultiLineString; use crate::extension::MultiPoint; @@ -19,7 +19,7 @@ use crate::extension::Point; use crate::extension::Polygon; use crate::extension::Rect; use crate::extension::WellKnownBinary; -use crate::prune::GeoDistanceBoundsPrune; +use crate::prune::GeoDistancePrune; use crate::scalar_fn::contains::GeoContains; use crate::scalar_fn::distance::GeoDistance; use crate::scalar_fn::intersects::GeoIntersects; @@ -66,9 +66,10 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(GeoDistance); session.scalar_fns().register(GeoIntersects); - // The bounding-box aggregate; self-declares as a per-chunk zone stat for geometry columns. - session.aggregate_fns().register(GeometryBounds); + // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for + // geometry columns. + session.aggregate_fns().register(GeometryAabb); - // Register the spatial pruning rule that uses that bounding box. - session.stats().register_rewrite(GeoDistanceBoundsPrune); + // Register the spatial pruning rule that uses that AABB. + session.stats().register_rewrite(GeoDistancePrune); } diff --git a/vortex-geo/src/prune.rs b/vortex-geo/src/prune.rs index 9960696c7f3..6b86c57c3e2 100644 --- a/vortex-geo/src/prune.rs +++ b/vortex-geo/src/prune.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Chunk pruning for spatial filters, using the per-chunk [`GeometryBounds`] bounding box. +//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned +//! bounding box (AABB). use geo::BoundingRect; use geo::Rect as GeoRect; @@ -32,13 +33,13 @@ use vortex_array::stats::stat; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::aggregate_fn::GeometryBounds; +use crate::aggregate_fn::GeometryAabb; use crate::extension::is_native_geometry; use crate::extension::single_geometry; use crate::scalar_fn::distance::GeoDistance; -/// Prunes chunks for `ST_Distance(geom, const) r` filters using the chunk's [`GeometryBounds`] -/// bounding box. Register it with `crate::initialize`; a chunk written without the `GeometryBounds` +/// Prunes chunks for `ST_Distance(geom, const) r` filters using the chunk's [`GeometryAabb`] +/// bounding box. Register it with `crate::initialize`; a chunk written without the `GeometryAabb` /// statistic is scanned rather than skipped. /// /// All four comparisons prune: `<= r` / `< r` skip a chunk whose box is wholly beyond `r` (box @@ -47,9 +48,9 @@ use crate::scalar_fn::distance::GeoDistance; /// `geometry_and_constant` + `distance_prune_proof` helpers; no new statistic or file-format change /// is needed. #[derive(Debug)] -pub struct GeoDistanceBoundsPrune; +pub struct GeoDistancePrune; -impl StatsRewriteRule for GeoDistanceBoundsPrune { +impl StatsRewriteRule for GeoDistancePrune { fn scalar_fn_id(&self) -> ScalarFnId { // The predicate root is the comparison, not `GeoDistance`, so key on `Binary`. Binary.id() @@ -60,7 +61,10 @@ impl StatsRewriteRule for GeoDistanceBoundsPrune { expr: &Expression, ctx: &StatsRewriteCtx<'_>, ) -> VortexResult> { - // Only the distance comparisons prune; `==` / `!=` and other operators fall through. + // Only the ordered comparisons prune today. `== r` could prune in the future (a chunk is + // provably empty when `r` lies outside its box's [min, max] distance interval), it's just + // not implemented. `!= r` cannot: pruning would need every row's distance to equal `r`, + // which an AABB can't prove. let op = *expr.as_::(); if !matches!( op, @@ -80,6 +84,9 @@ impl StatsRewriteRule for GeoDistanceBoundsPrune { let Some(radius) = expr.child(1).as_opt::() else { return Ok(None); }; + // Casts any primitive radius (integer literals included); it fails only for a null or + // non-primitive literal, where falling through means "scan the chunk", which is always + // sound. let Ok(radius) = f64::try_from(radius) else { return Ok(None); }; @@ -89,7 +96,7 @@ impl StatsRewriteRule for GeoDistanceBoundsPrune { return Ok(None); } - // Reduce `const` (any geometry type) to its bounding box. Every row sits in the chunk MBR + // Reduce `const` (any geometry type) to its AABB. Every row sits in the chunk AABB // and `const` in this box, so the box-to-box distance bounds the true distance soundly for // any geometry types. let mut exec = ctx.session().create_execution_ctx(); @@ -100,9 +107,9 @@ impl StatsRewriteRule for GeoDistanceBoundsPrune { } } -/// Shared bounds-pruning helper: split a symmetric geo predicate's operands into the scope-rooted +/// Shared AABB-pruning helper: split a symmetric geo predicate's operands into the scope-rooted /// geometry column and the constant's scalar, or `None` when the expression doesn't have that -/// shape or the geometry's dtype has no [`GeometryBounds`] support. Symmetric only - an asymmetric +/// shape or the geometry's dtype has no [`GeometryAabb`] support. Symmetric only - an asymmetric /// predicate that needs to know *which* operand is the column must recover the role separately. fn geometry_and_constant<'a>( expr: &'a Expression, @@ -119,7 +126,7 @@ fn geometry_and_constant<'a>( return Ok(None); }; - // A `GeometryBounds` stat reference only binds for dtypes it supports; anything else (e.g. a + // A `GeometryAabb` stat reference only binds for dtypes it supports; anything else (e.g. a // WKB column) must fall through to the scan. if !is_native_geometry(&ctx.return_dtype(geom)?) { return Ok(None); @@ -152,24 +159,24 @@ fn distance_prune_proof( } // The stat is read through `ext_storage`/`get_item`, which propagate a missing stat (null) to // "keep the chunk". Compared squared to avoid a `sqrt`; all operands are `>= 0`. - let mbr = ext_storage(stat(geom.clone(), GeometryBounds.bind(EmptyOptions))); + let aabb = ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))); let r2 = lit(radius * radius); Some(match op { // Beyond the threshold: even the nearest the box can be exceeds `r`. - Operator::Lte => gt(min_dist_sq(&mbr, query), r2), - Operator::Lt => gt_eq(min_dist_sq(&mbr, query), r2), + Operator::Lte => gt(min_dist_sq(&aabb, query), r2), + Operator::Lt => gt_eq(min_dist_sq(&aabb, query), r2), // Within the threshold: even the farthest the box can be is below `r`. - Operator::Gte => lt(max_dist_sq(&mbr, query), r2), - Operator::Gt => lt_eq(max_dist_sq(&mbr, query), r2), + Operator::Gte => lt(max_dist_sq(&aabb, query), r2), + Operator::Gt => lt_eq(max_dist_sq(&aabb, query), r2), _ => return None, }) } -/// Squared minimum distance between the chunk box `mbr` and the query box - a lower bound on every +/// Squared minimum distance between the chunk box `aabb` and the query box - a lower bound on every /// row's distance. `dx^2 + dy^2`, each axis gap clamped at zero (zero when the intervals overlap). -fn min_dist_sq(mbr: &Expression, query: GeoRect) -> Expression { - let field = |name: &str| get_item(name, mbr.clone()); - // max(0, q_lo - mbr_hi, mbr_lo - q_hi): positive only when the intervals are separated. +fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // max(0, q_lo - aabb_hi, aabb_lo - q_hi): positive only when the intervals are separated. let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { maximum( lit(0.0), @@ -184,12 +191,12 @@ fn min_dist_sq(mbr: &Expression, query: GeoRect) -> Expression { checked_add(square(dx), square(dy)) } -/// Squared maximum distance between the chunk box `mbr` and the query box - an upper bound on every +/// Squared maximum distance between the chunk box `aabb` and the query box - an upper bound on every /// row's distance. `Dx^2 + Dy^2`, each axis span the full extent of the two intervals' union. -fn max_dist_sq(mbr: &Expression, query: GeoRect) -> Expression { - let field = |name: &str| get_item(name, mbr.clone()); - // max(q_hi, mbr_hi) - min(q_lo, mbr_lo): the farthest two points of the boxes can be on an axis. - // The (nullable) MBR field is passed as the second arg so `case_when`'s else branch carries the +fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // max(q_hi, aabb_hi) - min(q_lo, aabb_lo): the farthest two points of the boxes can be on an axis. + // The (nullable) AABB field is passed as the second arg so `case_when`'s else branch carries the // nullability - a missing stat then propagates null through to "keep the chunk". let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { binop( @@ -257,11 +264,12 @@ mod tests { use vortex_error::VortexResult; use vortex_layout::layouts::zoned::zone_map::ZoneMap; - use super::GeoDistanceBoundsPrune; - use crate::aggregate_fn::GeometryBounds; + use super::GeoDistancePrune; + use crate::aggregate_fn::GeometryAabb; use crate::extension::GeoMetadata; use crate::extension::Rect; use crate::scalar_fn::distance::GeoDistance; + use crate::test_harness::geo_session; use crate::test_harness::point_column; /// Run the rule against `GeoDistance(root, origin) radius`, operands swapped when @@ -271,8 +279,7 @@ mod tests { geom_first: bool, radius: f64, ) -> VortexResult> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); @@ -285,7 +292,7 @@ mod tests { let distance = GeoDistance.new_expr(EmptyOptions, operands); let predicate = Binary.new_expr(operator, [distance, lit(radius)]); - GeoDistanceBoundsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + GeoDistancePrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) } /// All four distance comparisons prune (`<=`/`<` via min-distance, `>=`/`>` via max-distance); @@ -326,12 +333,11 @@ mod tests { Ok(()) } - /// A scope dtype without `GeometryBounds` support gets no proof - the stat reference would + /// A scope dtype without `GeometryAabb` support gets no proof - the stat reference would /// fail to bind at prune time. #[test] fn unsupported_scope_is_not_pruned() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let scope = DType::Primitive(PType::F64, Nullability::NonNullable); @@ -340,34 +346,32 @@ mod tests { let predicate = lt_eq(distance, lit(0.5f64)); let ctx = StatsRewriteCtx::new(&session, &scope); - assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); Ok(()) } /// A comparison that does not wrap `GeoDistance` is left untouched. #[test] fn ignores_non_distance_comparison() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); let predicate = lt_eq(lit(1.0f64), lit(2.0f64)); let ctx = StatsRewriteCtx::new(&session, &scope); - assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); Ok(()) } /// End-to-end over a hand-built zone map: the far chunk is skipped, the near one kept. #[test] fn prunes_far_chunk_keeps_near() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - // Two chunks' MBRs, stored as the native `geoarrow.box` stat column with default + // Two chunks' AABBs, stored as the native `geoarrow.box` stat column with default // (unreferenced) metadata to match the aggregate's return dtype: chunk 0 near the origin // (0,0..1,1), chunk 1 far away (100,100..101,101). let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); @@ -385,10 +389,10 @@ mod tests { .into_array(); let box_dtype = ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 2)?; + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 2)?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); @@ -407,14 +411,13 @@ mod tests { /// neither axis alone exceeds `r` - the case a per-axis box-overlap test would wrongly keep. #[test] fn prunes_diagonally_distant_chunk() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - // One chunk, MBR (0.8,0.8)..(0.9,0.9): each axis is only 0.8 from the origin (<= r = 1), but + // One chunk, AABB (0.8,0.8)..(0.9,0.9): each axis is only 0.8 from the origin (<= r = 1), but // the near corner is sqrt(0.8^2 + 0.8^2) ~= 1.13 away (> 1), so no point in the box is within 1. let ord = |a: f64| PrimitiveArray::from_iter([a]).into_array(); let boxes = StructArray::try_new( @@ -426,10 +429,10 @@ mod tests { .into_array(); let box_dtype = ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 1)?; + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 1)?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); @@ -448,12 +451,11 @@ mod tests { Ok(()) } - /// Backward compat: a zone map written without the `GeometryBounds` stat (an older file) keeps + /// Backward compat: a zone map written without the `GeometryAabb` stat (an older file) keeps /// every zone - the missing stat binds to null and `null_as_false` retains the zone. #[test] - fn missing_bounds_stat_keeps_all_zones() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); @@ -480,14 +482,13 @@ mod tests { /// satisfy `>= r`) via the box max-distance, while a chunk beyond `r` is kept. #[test] fn prunes_within_chunk_for_far_filter() -> VortexResult<()> { - let session = vortex_array::array_session(); - crate::initialize(&session); + let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - // Chunk 0 (MBR 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; + // Chunk 0 (AABB 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; // chunk 1 (100,100..101,101) is entirely beyond it. let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); let boxes = StructArray::try_new( @@ -504,10 +505,10 @@ mod tests { .into_array(); let box_dtype = ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let mbrs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); + let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 2)?; + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 2)?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 34f9862f4b9..79dc2d3361d 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -17,6 +17,7 @@ use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_session::VortexSession; use crate::extension::GeoMetadata; use crate::extension::LineString; @@ -36,6 +37,13 @@ use crate::extension::multipoint_storage_dtype; use crate::extension::multipolygon_storage_dtype; use crate::extension::polygon_storage_dtype; +/// A fresh session with the geospatial types, functions, and pruning rules registered. +pub(crate) fn geo_session() -> VortexSession { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +} + /// The WGS 84 (`EPSG:4326`) metadata tagged onto test geometry columns. fn wgs84() -> GeoMetadata { GeoMetadata { diff --git a/vortex-geo/src/tests/mod.rs b/vortex-geo/src/tests/mod.rs index 3708d5ad1fb..3ee1714ce47 100644 --- a/vortex-geo/src/tests/mod.rs +++ b/vortex-geo/src/tests/mod.rs @@ -17,8 +17,4 @@ use std::sync::LazyLock; use vortex_session::VortexSession; /// A session with the geospatial types and functions registered. -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(crate::test_harness::geo_session);