feat(spatial): write geometry bounding boxes to file statistics

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<AggregateFnRef> {
         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<AggregateFnRef> {
+        None
+    }
 }
 
 impl std::fmt::Debug for dyn AggregateFnPlugin {
@@ -69,4 +75,8 @@ impl<V: AggregateFnVTable> AggregateFnPlugin for V {
     fn zone_stat_default(&self, input_dtype: &DType) -> Option<AggregateFnRef> {
         AggregateFnVTable::zone_stat_default(self, input_dtype)
     }
+
+    fn file_stat_default(&self, input_dtype: &DType) -> Option<AggregateFnRef> {
+        AggregateFnVTable::file_stat_default(self, input_dtype)
+    }
 }

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<AggregateFnRef> {
+        self.registry.read(|registry| {
+            let mut fns: Vec<AggregateFnRef> = 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

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<AggregateFnRef> {
+        None
+    }
+
     /// DType of the intermediate partial accumulator state.
     ///
     /// Use a struct dtype when multiple fields are needed

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.

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)?;

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"] }
 

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:
 

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<AggregateFnRef> {
+        is_native_geometry(input_dtype).then(|| self.bind(EmptyOptions))
+    }
+
     fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
         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(())
     }
 }

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.