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 <zyu379@wisc.edu>

Author: HarukiMoriarty

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<AggregateFnRef>;
 
-    /// 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<AggregateFnRef> {
         None
     }

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

vortex-geo/src/aggregate_fn/bounds.rs vortex-geo/src/aggregate_fn/aabb.rsvortex-geo/src/aggregate_fn/bounds.rs renamed to vortex-geo/src/aggregate_fn/aabb.rs

@@ -3,6 +3,6 @@
 
 //! Aggregate functions over geometry columns.
 
-mod bounds;
+mod aabb;
 
-pub use bounds::*;
+pub use aabb::*;

vortex-geo/src/aggregate_fn/mod.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<xmin, ymin[, ..], xmax, ymax[, ..]>`.
+/// An axis-aligned bounding box (`geoarrow.box`), stored as `Struct<xmin, ymin[, ..], xmax, ymax[, ..]>`.
 // 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;

vortex-geo/src/extension/rect.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);
 }

vortex-geo/src/lib.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 {

vortex-geo/src/prune.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<VortexSession> = LazyLock::new(|| {
-    let session = vortex_array::array_session();
-    crate::initialize(&session);
-    session
-});
+static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::test_harness::geo_session);