From a5e57e7026987ed0104d4a6d67263d6043c9d3c1 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 26 Oct 2026 14:39:20 -0500 Subject: [PATCH] feat(spatial): write geometry bounding boxes to file statistics Signed-off-by: Nemo Yu --- Cargo.lock | 2 + vortex-array/src/aggregate_fn/plugin.rs | 10 ++ vortex-array/src/aggregate_fn/session.rs | 15 ++ vortex-array/src/aggregate_fn/vtable.rs | 6 + vortex-file/src/writer.rs | 2 +- vortex-layout/src/layouts/file_stats.rs | 44 +++-- vortex-spatial/Cargo.toml | 2 + vortex-spatial/README.md | 6 +- vortex-spatial/src/aggregate_fn/aabb.rs | 31 +++- vortex-spatial/src/lib.rs | 4 +- vortex-spatial/src/tests/file_stats.rs | 194 +++++++++++++++++++++++ vortex-spatial/src/tests/mod.rs | 3 +- 12 files changed, 293 insertions(+), 26 deletions(-) create mode 100644 vortex-spatial/src/tests/file_stats.rs diff --git a/Cargo.lock b/Cargo.lock index 8b18c9c2113..890d5fc7366 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12505,6 +12505,8 @@ dependencies = [ "vortex-edition", "vortex-error", "vortex-expr", + "vortex-file", + "vortex-io", "vortex-layout", "vortex-mask", "vortex-session", diff --git a/vortex-array/src/aggregate_fn/plugin.rs b/vortex-array/src/aggregate_fn/plugin.rs index baf6f303466..27d675a8480 100644 --- a/vortex-array/src/aggregate_fn/plugin.rs +++ b/vortex-array/src/aggregate_fn/plugin.rs @@ -42,6 +42,12 @@ pub trait AggregateFnPlugin: 'static + Send + Sync { fn zone_stat_default(&self, _input_dtype: &DType) -> Option { None } + + /// The default file statistic for a column of `input_dtype`, or `None` if this plugin does + /// not provide one. File writers discover these through the session's aggregate registry. + fn file_stat_default(&self, _input_dtype: &DType) -> Option { + None + } } impl std::fmt::Debug for dyn AggregateFnPlugin { @@ -69,4 +75,8 @@ impl AggregateFnPlugin for V { fn zone_stat_default(&self, input_dtype: &DType) -> Option { AggregateFnVTable::zone_stat_default(self, input_dtype) } + + fn file_stat_default(&self, input_dtype: &DType) -> Option { + AggregateFnVTable::file_stat_default(self, input_dtype) + } } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 9726b5e2ff4..d2305ddf0f9 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -192,6 +192,21 @@ impl AggregateFnSession { }) } + /// The default file statistics for a column of `input_dtype`, collected from every + /// registered aggregate's `file_stat_default`. + /// + /// This scans the plugin registry once per column when a file writer is initialized. + pub fn file_stat_defaults(&self, input_dtype: &DType) -> Vec { + self.registry.read(|registry| { + let mut fns: Vec = registry + .values() + .filter_map(|plugin| plugin.file_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 ba236ab1d5f..e057b8385d3 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -175,6 +175,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { None } + /// If this aggregate should be computed as a default file statistic for `input_dtype`, return + /// the bound aggregate to store. Default: not a file-statistics default. + fn file_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-file/src/writer.rs b/vortex-file/src/writer.rs index 14b43531695..6a81856ff04 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -159,7 +159,7 @@ impl VortexWriteOptions { /// Configure which statistics to compute at the file level. /// /// Pass an empty vector to omit file-level statistics. If left unset, the writer chooses - /// pruning aggregates from the file dtype. + /// pruning aggregates from the file dtype, including defaults registered by aggregate plugins. /// /// Like zone maps, file statistics only record aggregates the enabled editions include: the /// write fails if this asks for one they don't, while the defaults leave such aggregates out. diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 4e5a175b8b2..a0f7f7fee3c 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -25,6 +25,7 @@ use vortex_array::aggregate_fn::fns::max::Max; 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::session::AggregateFnSessionExt; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; @@ -49,8 +50,8 @@ use crate::sequence::SequentialStreamExt; /// Accumulates file statistics over `stream` as it passes through. /// -/// `stats` are the aggregates to record for every field, or `None` for each field's default (see -/// `default_pruning_aggregate_fns`). Only aggregates `ctx` allows are recorded in the nested +/// `stats` are the aggregates to record for every field, or `None` for each field's built-in and +/// session-registered defaults. Only aggregates `ctx` allows are recorded in the nested /// statistics: an explicitly requested aggregate it forbids fails, like it does for zone maps, /// while forbidden defaults are left out. /// @@ -372,7 +373,19 @@ enum StatsNode { } impl StatsNode { - fn build(dtype: &DType, stats: Option<&[AggregateFnRef]>) -> Self { + fn build(dtype: &DType, stats: Option<&[AggregateFnRef]>, session: &VortexSession) -> Self { + let leaf = || { + let mut defaults; + let stats = match stats { + Some(stats) => stats, + None => { + defaults = default_pruning_aggregate_fns(dtype); + defaults.extend(session.aggregate_fns().file_stat_defaults(dtype)); + defaults.as_slice() + } + }; + Self::Leaf(StatsAccumulator::new(dtype, Some(stats))) + }; match dtype.as_struct_fields_opt() { Some(struct_fields) => { let nullable = dtype.nullability() == Nullability::Nullable; @@ -386,25 +399,21 @@ impl StatsNode { entries: postorder_stats_layout(&field_dtype).len(), } } else { - Self::build(&field_dtype, stats) + Self::build(&field_dtype, stats, session) }; (Field::Name(name.clone()), child) }) .collect(); // Only reached through non-nullable ancestors, so the struct's own validity is // its full logical validity. - let own = if nullable { - Self::Leaf(StatsAccumulator::new(dtype, stats)) - } else { - Self::Skipped - }; + let own = if nullable { leaf() } else { Self::Skipped }; Self::Container { children, own: Box::new(own), } } None if !supports_file_stats(dtype) => Self::Skipped, - None => Self::Leaf(StatsAccumulator::new(dtype, stats)), + None => leaf(), } } @@ -512,10 +521,15 @@ impl FileStatsAccumulator { } } - let root = Arc::new(Mutex::new(StatsNode::build(dtype, stats.as_deref()))); + let root = Arc::new(Mutex::new(StatsNode::build( + dtype, + stats.as_deref(), + session, + ))); // Readers that predate nested stats reject a file whose legacy layout doesn't have one entry // per top-level field, so entries that aren't recorded are still written, empty. + // Plugin defaults are only accumulated in `root`: the legacy format cannot store them. let legacy = Arc::new(Mutex::new(match dtype.as_struct_fields_opt() { Some(struct_fields) => { // The fields of a nullable root aren't recorded, like `StatsNode::Unrecorded`. @@ -897,7 +911,7 @@ mod tests { StructArray::new(FieldNames::from(["a"]), [struct_a], 3, root_validity).into_array(); let requested = [agg(Stat::NullCount), agg(Stat::Min), agg(Stat::Max)]; - let mut node = StatsNode::build(root.dtype(), Some(&requested)); + let mut node = StatsNode::build(root.dtype(), Some(&requested), &SESSION); node.push_chunk(&root, &mut ctx)?; let stats_sets = node_stats_sets(&node)?; @@ -925,7 +939,7 @@ mod tests { StructArray::new(FieldNames::from(["s"]), [s], 3, Validity::NonNullable).into_array(); let requested = [agg(Stat::NullCount), agg(Stat::Min), agg(Stat::Max)]; - let mut node = StatsNode::build(root.dtype(), Some(&requested)); + let mut node = StatsNode::build(root.dtype(), Some(&requested), &SESSION); node.push_chunk(&root, &mut ctx)?; let mut aggregates = Vec::new(); @@ -953,7 +967,7 @@ mod tests { StructArray::new(FieldNames::from(["s"]), [s], 3, Validity::NonNullable).into_array(); let requested = [agg(Stat::NullCount), agg(Stat::Max)]; - let mut node = StatsNode::build(root.dtype(), Some(&requested)); + let mut node = StatsNode::build(root.dtype(), Some(&requested), &SESSION); node.push_chunk(&root, &mut ctx)?; let stats_sets = node_stats_sets(&node)?; @@ -981,7 +995,7 @@ mod tests { .into_array(); let requested = [agg(Stat::NullCount), agg(Stat::Min), agg(Stat::Max)]; - let mut node = StatsNode::build(outer.dtype(), Some(&requested)); + let mut node = StatsNode::build(outer.dtype(), Some(&requested), &SESSION); node.push_chunk(&outer, &mut ctx)?; let stats_sets = node_stats_sets(&node)?; diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index b48831cf879..bfb82b4e652 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -39,6 +39,8 @@ divan = { workspace = true } mimalloc = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-file = { workspace = true } +vortex-io = { workspace = true } vortex-layout = { workspace = true } vortex-spatial = { path = ".", features = ["_test-harness"] } diff --git a/vortex-spatial/README.md b/vortex-spatial/README.md index c3af7eec152..004183b0285 100644 --- a/vortex-spatial/README.md +++ b/vortex-spatial/README.md @@ -56,9 +56,11 @@ Most functions that need a general geometry algorithm use the row-function adapt Avoid per-row `scalar_at` or validity lookups in either path. Materialize coordinates, offsets, and validity once for the batch. -## Chunk pruning +## File statistics and chunk pruning -`GeometryAabb` computes the two-dimensional axis-aligned bounding box (AABB) of a geometry column. The writer stores it as the default zone statistic for native geometry columns. +`GeometryAabb` computes the two-dimensional axis-aligned bounding box (AABB) of a geometry column. After spatial initialization, the writer stores it by default in file statistics and zone maps for native geometry columns. + +File statistics accumulate the box across incoming batches and store the result in the footer, independently of the layout's zone maps. Consumers can select the column with `FileStatistics::get_by_path` and request `GeometryAabb.bind(EmptyOptions)` from its `AggregateStats`. The value is a native `geoarrow.box`; null or empty geometries contribute no extent. An explicit `with_file_statistics` configuration overrides these defaults. The current rules can prune: diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index 3ef24a619a7..7ab80318484 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/vortex-spatial/src/aggregate_fn/aabb.rs @@ -34,8 +34,8 @@ use crate::extension::flatten_coordinates; use crate::extension::is_native_geometry; /// 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`). +/// `geoarrow.box`, the spatial analogue of min/max. Also the default file and zone statistic for +/// such columns. #[derive(Clone, Debug)] pub struct GeometryAabb; @@ -150,6 +150,10 @@ impl AggregateFnVTable for GeometryAabb { is_native_geometry(input_dtype).then(|| self.bind(EmptyOptions)) } + fn file_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) } @@ -266,6 +270,7 @@ mod tests { use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateDTypes; use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; @@ -488,13 +493,19 @@ mod tests { 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. + /// Initializing spatial opts native geometry columns into file and zone statistics. #[test] - fn registered_as_geometry_zone_default() -> VortexResult<()> { + fn registered_as_geometry_stat_default() -> VortexResult<()> { + let uninitialized = VortexSession::empty(); let session = spatial_session(); for column in every_native_column(&[(0.0, 0.0), (1.0, 1.0)])? { + assert!( + uninitialized + .aggregate_fns() + .file_stat_defaults(column.dtype()) + .is_empty() + ); assert!( !session .aggregate_fns() @@ -503,6 +514,10 @@ mod tests { "a geometry zone-stat default should be discovered for {}", column.dtype() ); + assert_eq!( + session.aggregate_fns().file_stat_defaults(column.dtype()), + vec![GeometryAabb.bind(EmptyOptions)] + ); } let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert!( @@ -512,6 +527,12 @@ mod tests { .is_empty(), "no geometry zone-stat default should apply to numeric columns" ); + assert!( + session + .aggregate_fns() + .file_stat_defaults(&i32_dtype) + .is_empty() + ); Ok(()) } } diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index ecb3e1e55fd..76b9da1e73a 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -84,8 +84,8 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(SpatialMakeLine); session.scalar_fns().register(SpatialLength); - // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for - // geometry columns. + // The axis-aligned bounding-box (AABB) aggregate; self-declares as a file and zone statistic + // for geometry columns. session.aggregate_fns().register(GeometryAabb); // Register the spatial pruning rules that use that AABB. diff --git a/vortex-spatial/src/tests/file_stats.rs b/vortex-spatial/src/tests/file_stats.rs new file mode 100644 index 00000000000..ddb0317dcf0 --- /dev/null +++ b/vortex-spatial/src/tests/file_stats.rs @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::fns::null_count::NullCount; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::FieldPath; +use vortex_array::dtype::Nullability; +use vortex_array::expr::stats::Precision; +use vortex_array::flatbuffers::WriteFlatBufferExt; +use vortex_array::iter::ArrayIteratorAdapter; +use vortex_array::iter::ArrayIteratorExt; +use vortex_array::scalar::Scalar; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::EDITION_FAMILIES; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::runtime::Handle; +use vortex_io::runtime::single::block_on; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_session::VortexSession; + +use crate::aggregate_fn::GeometryAabb; +use crate::extension::Rect; +use crate::extension::SpatialMetadata; +use crate::extension::box_storage_dtype; +use crate::extension::coordinate::Dimension; +use crate::test_harness::nullable_point_column; + +fn file_session(handle: Handle) -> VortexResult { + let session = VortexSession::empty().with_handle(handle); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + vortex_file::register_default_encodings(&session); + Ok(session) +} + +fn geometry_batch(points: Vec>, nested: bool) -> VortexResult { + let batch = StructArray::from_fields(&[("geom", nullable_point_column(points)?)])?.into_array(); + if nested { + Ok(StructArray::from_fields(&[("location", batch)])?.into_array()) + } else { + Ok(batch) + } +} + +#[rstest] +fn default_file_bbox_spans_batches( + #[values(false, true)] nested: bool, + #[values(false, true)] without_zones: bool, +) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + crate::initialize(&session); + + let batches = [ + geometry_batch(vec![Some((10.0, 20.0)), None], nested)?, + geometry_batch(vec![None], nested)?, + geometry_batch(vec![Some((-10.0, 80.0)), Some((30.0, 40.0))], nested)?, + ]; + let dtype = batches[0].dtype().clone(); + let stream = + ArrayIteratorAdapter::new(dtype, batches.into_iter().map(Ok)).into_array_stream(); + let mut options = session.write_options(); + if without_zones { + options = options.with_strategy(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))); + } + let mut buffer = ByteBufferMut::empty(); + let summary = options.write(&mut buffer, stream).await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + let path = if nested { + FieldPath::from_name("location").push("geom") + } else { + FieldPath::from_name("geom") + }; + let expected = Scalar::extension::( + SpatialMetadata::default(), + Scalar::struct_( + box_storage_dtype(Dimension::Xy, Nullability::Nullable), + vec![(-10.0).into(), 20.0.into(), 30.0.into(), 80.0.into()], + ), + ); + for stats in [summary.footer().statistics(), file.file_stats()] { + let (stats, _) = stats + .and_then(|stats| stats.get_by_path(&path)) + .ok_or_else(|| vortex_err!("missing geometry file statistics"))?; + assert_eq!( + stats.get(&GeometryAabb.bind(EmptyOptions)), + Precision::exact(expected.clone()) + ); + assert_eq!( + stats.get(&NullCount.bind(EmptyOptions)), + Precision::exact(Scalar::from(2u64)) + ); + } + Ok(()) + }) +} + +#[rstest] +fn explicit_file_statistics_override_spatial_default( + #[values(false, true)] disabled: bool, +) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + crate::initialize(&session); + let batch = geometry_batch(vec![Some((10.0, 20.0)), None], false)?; + let requested = if disabled { + vec![] + } else { + vec![NullCount.bind(EmptyOptions)] + }; + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_file_statistics(requested) + .write(&mut buffer, batch.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + if disabled { + assert!(file.file_stats().is_none()); + } else { + let (stats, _) = file + .file_stats() + .and_then(|stats| stats.get_by_path(&FieldPath::from_name("geom"))) + .ok_or_else(|| vortex_err!("missing geometry file statistics"))?; + assert!(stats.get(&GeometryAabb.bind(EmptyOptions)).is_absent()); + assert_eq!( + stats.get(&NullCount.bind(EmptyOptions)), + Precision::exact(Scalar::from(1u64)) + ); + } + Ok(()) + }) +} + +#[test] +fn spatial_initialization_preserves_ordinary_file_statistics() -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + let batch = StructArray::from_fields(&[( + "value", + PrimitiveArray::from_option_iter([Some(3i32), None, Some(-2)]).into_array(), + )])? + .into_array(); + let mut statistics = Vec::new(); + for spatial in [false, true] { + if spatial { + crate::initialize(&session); + } + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buffer, batch.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + statistics.push( + file.file_stats() + .ok_or_else(|| vortex_err!("missing file statistics"))? + .write_flatbuffer_bytes()?, + ); + } + assert_eq!(statistics[0], statistics[1]); + Ok(()) + }) +} diff --git a/vortex-spatial/src/tests/mod.rs b/vortex-spatial/src/tests/mod.rs index c83a30d476d..fa27ab81a20 100644 --- a/vortex-spatial/src/tests/mod.rs +++ b/vortex-spatial/src/tests/mod.rs @@ -1,9 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Arrow interop tests for the spatial extension types, exercising the session wiring set up +//! Interop tests for the spatial extension types, exercising the session wiring set up //! by [`crate::initialize`]. +mod file_stats; mod linestring; mod multilinestring; mod multipoint;