Make stats binders immutable

Signed-off-by: "Nicholas Gates" <nick@nickgates.com>

Author: gatesn

vortex-array/src/expr/pruning/pruning_expr.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::cell::RefCell;
 use std::iter;
 
 use itertools::Itertools;
@@ -70,14 +71,14 @@ pub fn checked_pruning_expr(
         return Ok(None);
     };
 
-    let mut binder = RequiredStatsBinder {
+    let binder = RequiredStatsBinder {
         scope,
         available_stats,
-        required_stats: Relation::new(),
-        bound_stats: Vec::new(),
+        required_stats: RefCell::new(Relation::new()),
+        bound_stats: RefCell::new(Vec::new()),
     };
-    let lowered = bind_stats(predicate, &mut binder)?;
-    let required_stats = filter_required_stats(&lowered, binder.required_stats);
+    let lowered = bind_stats(predicate, &binder)?;
+    let required_stats = filter_required_stats(&lowered, binder.required_stats.into_inner());
     // If no stats-table fields remain, only a constant `true` proof can prune.
     // `false`, `null`, and non-constant expressions cannot justify building a
     // stats-table pruning expression.
@@ -91,8 +92,8 @@ pub fn checked_pruning_expr(
 struct RequiredStatsBinder<'a> {
     scope: &'a DType,
     available_stats: &'a FieldPathSet,
-    required_stats: RequiredStats,
-    bound_stats: Vec<(FieldName, DType)>,
+    required_stats: RefCell<RequiredStats>,
+    bound_stats: RefCell<Vec<(FieldName, DType)>>,
 }
 
 impl StatBinder for RequiredStatsBinder<'_> {
@@ -102,13 +103,13 @@ impl StatBinder for RequiredStatsBinder<'_> {
 
     fn bound_scope(&self) -> DType {
         DType::Struct(
-            StructFields::from_iter(self.bound_stats.iter().cloned()),
+            StructFields::from_iter(self.bound_stats.borrow().iter().cloned()),
             Nullability::NonNullable,
         )
     }
 
     fn bind_stat(
-        &mut self,
+        &self,
         input: &Expression,
         stat: Stat,
         stat_dtype: &DType,
@@ -129,21 +130,20 @@ impl StatBinder for RequiredStatsBinder<'_> {
         }
 
         let stat_field_name = field_path_stat_field_name(&field_path, stat);
-        if self
-            .bound_stats
+        let mut bound_stats = self.bound_stats.borrow_mut();
+        if bound_stats
             .iter()
             .all(|(field_name, _)| field_name != stat_field_name)
         {
-            self.bound_stats
-                .push((stat_field_name.clone(), stat_dtype.clone()));
+            bound_stats.push((stat_field_name.clone(), stat_dtype.clone()));
         }
 
-        self.required_stats.insert(field_path, stat);
+        self.required_stats.borrow_mut().insert(field_path, stat);
         Ok(Some(get_item(stat_field_name, root())))
     }
 
     fn bind_aggregate(
-        &mut self,
+        &self,
         input: &Expression,
         aggregate_fn: &AggregateFnRef,
         stat_dtype: &DType,

vortex-array/src/stats/bind.rs

@@ -47,7 +47,7 @@ pub trait StatBinder {
     /// then call [`Self::missing_stat`] with the dtype expected from the
     /// original `vortex.stat` expression.
     fn bind_stat(
-        &mut self,
+        &self,
         input: &Expression,
         stat: Stat,
         stat_dtype: &DType,
@@ -59,7 +59,7 @@ pub trait StatBinder {
     /// [`Stat`] slots. Binders that store richer aggregate stats can override
     /// this method without extending the generic stats binding walker.
     fn bind_aggregate(
-        &mut self,
+        &self,
         input: &Expression,
         aggregate_fn: &AggregateFnRef,
         stat_dtype: &DType,
@@ -68,11 +68,7 @@ pub trait StatBinder {
     }
 
     /// Bind one of the legacy stat slots for `input`.
-    fn bind_legacy_stat(
-        &mut self,
-        input: &Expression,
-        stat: Stat,
-    ) -> VortexResult<Option<Expression>> {
+    fn bind_legacy_stat(&self, input: &Expression, stat: Stat) -> VortexResult<Option<Expression>> {
         let input_dtype = input.return_dtype(self.scope())?;
         let Some(stat_dtype) = stat.dtype(&input_dtype) else {
             return Ok(None);
@@ -84,7 +80,7 @@ pub trait StatBinder {
     ///
     /// The default is a nullable null literal, which preserves three-valued
     /// pruning semantics for stats-table execution.
-    fn missing_stat(&mut self, dtype: DType) -> VortexResult<Expression> {
+    fn missing_stat(&self, dtype: DType) -> VortexResult<Expression> {
         Ok(null_expr(dtype))
     }
 }
@@ -94,7 +90,10 @@ pub trait StatBinder {
 /// The predicate is usually the output of a stats rewrite rule. Rewrite rules
 /// are responsible for expressing stat semantics; binding maps aggregate-backed
 /// stat requests to the concrete stats representation supported by the binder.
-pub fn bind_stats(predicate: Expression, binder: &mut impl StatBinder) -> VortexResult<Expression> {
+pub fn bind_stats<B: StatBinder + ?Sized>(
+    predicate: Expression,
+    binder: &B,
+) -> VortexResult<Expression> {
     let scope = binder.scope().clone();
     let lowered = predicate
         .transform_down(|expr| {
@@ -120,7 +119,7 @@ pub fn bind_stats(predicate: Expression, binder: &mut impl StatBinder) -> Vortex
 /// This is an opt-in helper for stats backends that materialize `NaNCount` and
 /// `NullCount`, but do not materialize aggregate boolean stats directly.
 pub fn bind_legacy_count_aggregate<B: StatBinder + ?Sized>(
-    binder: &mut B,
+    binder: &B,
     input: &Expression,
     aggregate_fn: &AggregateFnRef,
 ) -> VortexResult<Option<Expression>> {
@@ -157,7 +156,7 @@ pub fn bind_legacy_count_aggregate<B: StatBinder + ?Sized>(
 
 /// Bind an aggregate function that has a direct legacy [`Stat`] slot.
 pub fn bind_direct_aggregate_stat<B: StatBinder + ?Sized>(
-    binder: &mut B,
+    binder: &B,
     input: &Expression,
     aggregate_fn: &AggregateFnRef,
     stat_dtype: &DType,
@@ -174,7 +173,7 @@ pub fn bind_direct_aggregate_stat<B: StatBinder + ?Sized>(
 /// `NaNCount` and `NullCount`, then fall back to direct aggregate-to-stat
 /// mappings.
 pub fn bind_legacy_count_or_direct_aggregate<B: StatBinder + ?Sized>(
-    binder: &mut B,
+    binder: &B,
     input: &Expression,
     aggregate_fn: &AggregateFnRef,
     stat_dtype: &DType,
@@ -189,7 +188,7 @@ pub fn bind_legacy_count_or_direct_aggregate<B: StatBinder + ?Sized>(
 fn bind_stat_fn(
     expr: &Expression,
     scope: &DType,
-    binder: &mut impl StatBinder,
+    binder: &(impl StatBinder + ?Sized),
 ) -> VortexResult<Option<Expression>> {
     let options = expr.as_::<StatFn>();
     let aggregate_fn = options.aggregate_fn();
@@ -258,7 +257,7 @@ mod tests {
         }
 
         fn bind_stat(
-            &mut self,
+            &self,
             _input: &Expression,
             stat: Stat,
             _stat_dtype: &DType,
@@ -271,7 +270,7 @@ mod tests {
         }
 
         fn bind_aggregate(
-            &mut self,
+            &self,
             input: &Expression,
             aggregate_fn: &AggregateFnRef,
             stat_dtype: &DType,
@@ -282,33 +281,33 @@ mod tests {
 
     #[test]
     fn all_non_nan_binds_to_nan_count_zero() -> VortexResult<()> {
-        let mut binder = TestBinder::new(true);
+        let binder = TestBinder::new(true);
 
-        let bound = bind_stats(all_non_nan(col("f")), &mut binder)?;
+        let bound = bind_stats(all_non_nan(col("f")), &binder)?;
 
         assert_eq!(bound, eq(col("f_nan_count"), lit(0u64)));
         Ok(())
     }
 
     #[test]
     fn all_non_nan_lowers_to_null_when_nan_count_is_missing() -> VortexResult<()> {
-        let mut binder = TestBinder::new(false);
+        let binder = TestBinder::new(false);
 
-        let bound = bind_stats(all_non_nan(col("f")), &mut binder)?;
+        let bound = bind_stats(all_non_nan(col("f")), &binder)?;
 
         assert_eq!(bound, lit(Scalar::null(DType::Bool(Nullability::Nullable))));
         Ok(())
     }
 
     #[test]
     fn missing_stats_fold_when_kleene_semantics_allow_it() -> VortexResult<()> {
-        let mut binder = TestBinder::new(false);
+        let binder = TestBinder::new(false);
 
-        let bound = bind_stats(and(lit(false), all_non_nan(col("f"))), &mut binder)?;
+        let bound = bind_stats(and(lit(false), all_non_nan(col("f"))), &binder)?;
 
         assert_eq!(bound, lit(false));
 
-        let bound = bind_stats(or(lit(true), all_non_nan(col("f"))), &mut binder)?;
+        let bound = bind_stats(or(lit(true), all_non_nan(col("f"))), &binder)?;
 
         assert_eq!(bound, lit(true));
         Ok(())
@@ -328,7 +327,7 @@ mod tests {
             }
 
             fn bind_stat(
-                &mut self,
+                &self,
                 input: &Expression,
                 stat: Stat,
                 stat_dtype: &DType,
@@ -337,19 +336,19 @@ mod tests {
             }
         }
 
-        let mut binder = DefaultBinder(TestBinder::new(true));
+        let binder = DefaultBinder(TestBinder::new(true));
 
-        let bound = bind_stats(all_non_nan(col("f")), &mut binder)?;
+        let bound = bind_stats(all_non_nan(col("f")), &binder)?;
 
         assert_eq!(bound, lit(Scalar::null(DType::Bool(Nullability::Nullable))));
         Ok(())
     }
 
     #[test]
     fn unrelated_expressions_do_not_request_nan_count() -> VortexResult<()> {
-        let mut binder = TestBinder::new(false);
+        let binder = TestBinder::new(false);
 
-        let bound = bind_stats(is_null(col("f")), &mut binder)?;
+        let bound = bind_stats(is_null(col("f")), &binder)?;
 
         assert_eq!(bound, is_null(col("f")));
         Ok(())

vortex-file/src/v2/file_stats_reader.rs

@@ -121,8 +121,8 @@ impl FileStatsLayoutReader {
     }
 
     fn lower_stats(&self, predicate: Expression) -> VortexResult<Expression> {
-        let mut binder = FileStatsBinder { reader: self };
-        bind_stats(predicate, &mut binder)
+        let binder = FileStatsBinder { reader: self };
+        bind_stats(predicate, &binder)
     }
 
     pub fn file_stats(&self) -> &FileStatistics {
@@ -144,7 +144,7 @@ impl StatBinder for FileStatsBinder<'_> {
     }
 
     fn bind_stat(
-        &mut self,
+        &self,
         input: &Expression,
         stat: Stat,
         _stat_dtype: &DType,
@@ -156,7 +156,7 @@ impl StatBinder for FileStatsBinder<'_> {
     }
 
     fn bind_aggregate(
-        &mut self,
+        &self,
         input: &Expression,
         aggregate_fn: &AggregateFnRef,
         stat_dtype: &DType,

vortex-layout/src/layouts/zoned/zone_map.rs

@@ -122,8 +122,8 @@ impl ZoneMap {
     }
 
     fn lower_stats(&self, predicate: Expression) -> VortexResult<Expression> {
-        let mut binder = ZoneMapStatsBinder { zone_map: self };
-        bind_stats(predicate, &mut binder)
+        let binder = ZoneMapStatsBinder { zone_map: self };
+        bind_stats(predicate, &binder)
     }
 }
 
@@ -141,7 +141,7 @@ impl StatBinder for ZoneMapStatsBinder<'_> {
     }
 
     fn bind_stat(
-        &mut self,
+        &self,
         input: &Expression,
         stat: Stat,
         _stat_dtype: &DType,
@@ -161,7 +161,7 @@ impl StatBinder for ZoneMapStatsBinder<'_> {
     }
 
     fn bind_aggregate(
-        &mut self,
+        &self,
         input: &Expression,
         aggregate_fn: &AggregateFnRef,
         stat_dtype: &DType,