Consolidate into stat_expression function (#5366)

This removes the unused stat_field_path function, as well as improves
coverage for propagating other statistics.

Interestingly it has highlighted the difference between pruning stats
(min, max, is_null), and aggregate stats (sum, null_count, etc.). I have
some ideas for generalizing these concepts, but that's for another day.

---------

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/expr/analysis.rs

@@ -14,64 +14,7 @@ pub trait StatsCatalog {
     /// This is likely to be a column expression, or a literal.
     ///
     /// Returns `None` if the stat is not available for the field path.
-    fn stats_ref(&mut self, _field_path: &FieldPath, _stat: Stat) -> Option<Expression> {
+    fn stats_ref(&self, _field_path: &FieldPath, _stat: Stat) -> Option<Expression> {
         None
     }
 }
-
-/// This can be used by expression to plug into vortex expression analysis, such as
-/// pruning or expression simplification
-pub trait AnalysisExpr {
-    /// An expression over zone-statistics which implies all records in the zone evaluate to false.
-    ///
-    /// Given an expression, `e`, if `e.stat_falsification(..)` evaluates to true, it is guaranteed
-    /// that `e` evaluates to false on all records in the zone. However, the inverse is not
-    /// necessarily true: even if the falsification evaluates to false, `e` need not evaluate to
-    /// true on all records.
-    ///
-    /// The [`StatsCatalog`] can be used to constrain or rename stats used in the final expr.
-    ///
-    /// # Examples
-    ///
-    /// - An expression over one variable: `x > 0` is false for all records in a zone if the maximum
-    ///   value of the column `x` in that zone is less than or equal to zero: `max(x) <= 0`.
-    /// - An expression over two variables: `x > y` becomes `max(x) <= min(y)`.
-    /// - A conjunctive expression: `x > y AND z < x` becomes `max(x) <= min(y) OR min(z) >= max(x).
-    ///
-    /// Some expressions, in theory, have falsifications but this function does not support them
-    /// such as `x < (y < z)` or `x LIKE "needle%"`.
-    fn stat_falsification(&self, _catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        None
-    }
-
-    /// An expression for the upper non-null bound of this expression, if available.
-    ///
-    /// This function returns None if there is no upper bound or it is difficult to compute.
-    ///
-    /// The returned expression evaluates to null if the maximum value is unknown. In that case, you
-    /// _must not_ assume the array is empty _nor_ may you assume the array only contains non-null
-    /// values.
-    fn max(&self, _catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        None
-    }
-
-    /// An expression for the lower non-null bound of this expression, if available.
-    ///
-    /// See [AnalysisExpr::max] for important details.
-    fn min(&self, _catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        None
-    }
-
-    /// An expression for the NaN count for a column, if available.
-    ///
-    /// This method returns `None` if the NaNCount stat is unknown.
-    fn nan_count(&self, _catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        None
-    }
-
-    fn field_path(&self) -> Option<FieldPath> {
-        None
-    }
-
-    // TODO: add containment
-}

vortex-array/src/expr/expression.rs

@@ -7,13 +7,14 @@ use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
-use vortex_dtype::{DType, FieldPath};
+use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::Vector;
 
 use crate::ArrayRef;
 use crate::expr::display::DisplayTreeExpr;
 use crate::expr::{ChildName, ExprId, ExprVTable, ExpressionView, StatsCatalog, VTable};
+use crate::stats::Stat;
 
 /// A node in a Vortex expression tree.
 ///
@@ -164,39 +165,34 @@ impl Expression {
     ///
     /// Some expressions, in theory, have falsifications but this function does not support them
     /// such as `x < (y < z)` or `x LIKE "needle%"`.
-    pub fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<Expression> {
+    pub fn stat_falsification(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
         self.vtable.as_dyn().stat_falsification(self, catalog)
     }
 
-    /// An expression for the upper non-null bound of this expression, if available.
+    /// Returns an expression representing the zoned statistic for the given stat, if available.
     ///
-    /// This function returns None if there is no upper bound, or it is difficult to compute.
+    /// The [`StatsCatalog`] returns expressions that can be evaluated using the zone map as a
+    /// scope. Expressions can implement this function to propagate such statistics through the
+    /// expression tree. For example, the `a + 10` expression could propagate `min: min(a) + 10`.
     ///
-    /// The returned expression evaluates to null if the maximum value is unknown. In that case, you
-    /// _must not_ assume the array is empty _nor_ may you assume the array only contains non-null
-    /// values.
-    pub fn stat_max(&self, catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        self.vtable.as_dyn().stat_max(self, catalog)
+    /// NOTE(gatesn): we currently cannot represent statistics over nested fields. Please file an
+    /// issue to discuss a solution to this.
+    pub fn stat_expression(&self, stat: Stat, catalog: &dyn StatsCatalog) -> Option<Expression> {
+        self.vtable.as_dyn().stat_expression(self, stat, catalog)
     }
 
-    /// An expression for the lower non-null bound of this expression, if available.
+    /// Returns an expression representing the zoned maximum statistic, if available.
     ///
-    /// See [`Expression::stat_max`] for important details.
-    pub fn stat_min(&self, catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        self.vtable.as_dyn().stat_min(self, catalog)
+    /// See [`Self::stat_expression`] for details.
+    pub fn stat_min(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
+        self.stat_expression(Stat::Min, catalog)
     }
 
-    /// An expression for the NaN count for a column, if available.
+    /// Returns an expression representing the zoned maximum statistic, if available.
     ///
-    /// This method returns `None` if the NaNCount stat is unknown.
-    pub fn stat_nan_count(&self, catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        self.vtable.as_dyn().stat_nan_count(self, catalog)
-    }
-
-    // TODO(ngates): I'm not sure what this is really for? We need to clean up stats compute for
-    //  expressions.
-    pub fn stat_field_path(&self) -> Option<FieldPath> {
-        self.vtable.as_dyn().stat_field_path(self)
+    /// See [`Self::stat_expression`] for details.
+    pub fn stat_max(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
+        self.stat_expression(Stat::Max, catalog)
     }
 
     /// Format the expression as a compact string.

vortex-array/src/expr/exprs/between.rs

@@ -139,7 +139,7 @@ impl VTable for Between {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         expr.to_binary_expr().stat_falsification(catalog)
     }

vortex-array/src/expr/exprs/binary.rs

@@ -13,6 +13,7 @@ use crate::expr::expression::Expression;
 use crate::expr::exprs::literal::lit;
 use crate::expr::exprs::operators::Operator;
 use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::stats::Stat;
 use crate::{ArrayRef, compute};
 
 pub struct Binary;
@@ -104,7 +105,7 @@ impl VTable for Binary {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         // Wrap another predicate with an optional NaNCount check, if the stat is available.
         //
@@ -124,12 +125,12 @@ impl VTable for Binary {
             lhs: &Expression,
             rhs: &Expression,
             value_predicate: Expression,
-            catalog: &mut dyn StatsCatalog,
+            catalog: &dyn StatsCatalog,
         ) -> Expression {
             let nan_predicate = lhs
-                .stat_nan_count(catalog)
+                .stat_expression(Stat::NaNCount, catalog)
                 .into_iter()
-                .chain(rhs.stat_nan_count(catalog))
+                .chain(rhs.stat_expression(Stat::NaNCount, catalog))
                 .map(|nans| eq(nans, lit(0u64)))
                 .reduce(and);
 

vortex-array/src/expr/exprs/cast.rs

@@ -5,14 +5,15 @@ use std::fmt::Formatter;
 use std::ops::Deref;
 
 use prost::Message;
-use vortex_dtype::{DType, FieldPath};
+use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr as pb;
 
 use crate::ArrayRef;
 use crate::compute::cast as compute_cast;
 use crate::expr::expression::Expression;
 use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::stats::Stat;
 
 /// A cast expression that converts values to a target data type.
 pub struct Cast;
@@ -86,36 +87,36 @@ impl VTable for Cast {
         })
     }
 
-    fn stat_max(
+    fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        stat: Stat,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        expr.child(0)
-            .stat_max(catalog)
-            .map(|x| cast(x, expr.data().clone()))
-    }
-
-    fn stat_min(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
-    ) -> Option<Expression> {
-        expr.child(0)
-            .stat_min(catalog)
-            .map(|x| cast(x, expr.data().clone()))
-    }
-
-    fn stat_nan_count(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
-    ) -> Option<Expression> {
-        expr.child(0).stat_nan_count(catalog)
-    }
-
-    fn stat_field_path(&self, expr: &ExpressionView<Self>) -> Option<FieldPath> {
-        expr.children()[0].stat_field_path()
+        match stat {
+            Stat::IsConstant
+            | Stat::IsSorted
+            | Stat::IsStrictSorted
+            | Stat::NaNCount
+            | Stat::Sum
+            | Stat::UncompressedSizeInBytes => expr.child(0).stat_expression(stat, catalog),
+            Stat::Max | Stat::Min => {
+                // We cast min/max to the new type
+                expr.child(0)
+                    .stat_expression(stat, catalog)
+                    .map(|x| cast(x, expr.data().clone()))
+            }
+            Stat::NullCount => {
+                // if !expr.data().is_nullable() {
+                // NOTE(ngates): we should decide on the semantics here. In theory, the null
+                //  count of something cast to non-nullable will be zero. But if we return
+                //  that we know this to be zero, then a pruning predicate may eliminate data
+                //  that would otherwise have caused the cast to error.
+                // return Some(lit(0u64));
+                // }
+                None
+            }
+        }
     }
 }
 

vortex-array/src/expr/exprs/dynamic.rs

@@ -91,7 +91,7 @@ impl VTable for DynamicComparison {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<DynamicComparison>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         match expr.data().operator {
             Operator::Gt => Some(DynamicComparison.new_expr(

vortex-array/src/expr/exprs/get_item/mod.rs

@@ -121,34 +121,21 @@ impl VTable for GetItem {
         Ok(field)
     }
 
-    fn stat_max(
+    fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        stat: Stat,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), Stat::Max)
-    }
-
-    fn stat_min(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
-    ) -> Option<Expression> {
-        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), Stat::Min)
-    }
-
-    fn stat_nan_count(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
-    ) -> Option<Expression> {
-        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), Stat::NaNCount)
-    }
-
-    fn stat_field_path(&self, expr: &ExpressionView<Self>) -> Option<FieldPath> {
-        expr.children()[0]
-            .stat_field_path()
-            .map(|fp| fp.push(expr.data().clone()))
+        // TODO(ngates): I think we can do better here and support stats over nested fields.
+        //  It would be nice if delegating to our child would return a struct of statistics
+        //  matching the nested DType such that we can write:
+        //    `get_item(expr.child(0).stat_expression(...), expr.data().field_name())`
+
+        // TODO(ngates): this is a bug whereby we may return stats for a nested field of the same
+        //  name as a field in the root struct. This should be resolved with upcoming change to
+        //  falsify expressions, but for now I'm preserving the existing buggy behavior.
+        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), stat)
     }
 }
 

vortex-array/src/expr/exprs/is_null.rs

@@ -88,10 +88,9 @@ impl VTable for IsNull {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        let field_path = expr.children()[0].stat_field_path()?;
-        let null_count_expr = catalog.stats_ref(&field_path, Stat::NullCount)?;
+        let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?;
         Some(eq(null_count_expr, lit(0u64)))
     }
 }

vortex-array/src/expr/exprs/list_contains.rs

@@ -84,7 +84,7 @@ impl VTable for ListContains {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         // falsification(contains([1,2,5], x)) =>
         //   falsification(x != 1) and falsification(x != 2) and falsification(x != 5)