Stat expression

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,12 +7,13 @@ use std::fmt::{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 crate::ArrayRef;
 use crate::expr::display::DisplayTreeExpr;
 use crate::expr::{ChildName, ExprId, ExprVTable, ExpressionView, StatsCatalog, VTable};
+use crate::stats::Stat;
+use crate::ArrayRef;
 
 /// A node in a Vortex expression tree.
 ///
@@ -158,39 +159,17 @@ 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.
-    ///
-    /// 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.
-    pub fn stat_max(&self, catalog: &mut dyn StatsCatalog) -> Option<Expression> {
-        self.vtable.as_dyn().stat_max(self, catalog)
-    }
-
-    /// An expression for the lower non-null bound of this expression, if available.
+    /// Returns an expression representing the zoned statistic for the given stat, 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)
-    }
-
-    /// An expression for the NaN count for a column, 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)
+    /// 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`.
+    pub fn stat_expression(&self, stat: Stat, catalog: &dyn StatsCatalog) -> Option<Expression> {
+        self.vtable.as_dyn().stat_expression(self, stat, catalog)
     }
 
     /// Format the expression as a compact string.

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

@@ -6,15 +6,15 @@ use std::fmt::Formatter;
 use prost::Message;
 use vortex_dtype::DType;
 use vortex_dtype::DType::Bool;
-use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 use vortex_proto::expr as pb;
 
-use crate::ArrayRef;
-use crate::compute::{BetweenOptions, between as between_compute};
+use crate::compute::{between as between_compute, BetweenOptions};
 use crate::expr::expression::Expression;
 use crate::expr::exprs::binary::Binary;
 use crate::expr::exprs::operators::Operator;
 use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::ArrayRef;
 
 /// An optimized scalar expression to compute whether values fall between two bounds.
 ///
@@ -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

@@ -5,15 +5,15 @@ use std::fmt::Formatter;
 
 use prost::Message;
 use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 use vortex_proto::expr as pb;
 
 use crate::compute::{add, and_kleene, compare, div, mul, or_kleene, sub};
 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::{ArrayRef, compute};
+use crate::{compute, ArrayRef};
 
 pub struct Binary;
 
@@ -104,7 +104,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,7 +124,7 @@ 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)
@@ -510,7 +510,7 @@ mod tests {
     use super::{and, and_collect, and_collect_right, eq, gt, gt_eq, lt, lt_eq, not_eq, or};
     use crate::expr::exprs::get_item::col;
     use crate::expr::exprs::literal::lit;
-    use crate::expr::{Expression, test_harness};
+    use crate::expr::{test_harness, Expression};
 
     #[test]
     fn and_collect_left_assoc() {

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

@@ -6,13 +6,13 @@ use std::ops::Deref;
 
 use prost::Message;
 use vortex_dtype::{DType, FieldPath};
-use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
 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::ArrayRef;
 
 /// A cast expression that converts values to a target data type.
 pub struct Cast;
@@ -89,23 +89,23 @@ impl VTable for Cast {
     fn stat_max(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         expr.children()[0].stat_max(catalog)
     }
 
     fn stat_min(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         expr.children()[0].stat_min(catalog)
     }
 
     fn stat_nan_count(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         expr.children()[0].stat_nan_count(catalog)
     }
@@ -136,11 +136,11 @@ mod tests {
     use vortex_error::VortexUnwrap as _;
 
     use super::cast;
-    use crate::IntoArray;
     use crate::arrays::StructArray;
     use crate::expr::exprs::get_item::get_item;
     use crate::expr::exprs::root::root;
-    use crate::expr::{Expression, test_harness};
+    use crate::expr::{test_harness, Expression};
+    use crate::IntoArray;
 
     #[test]
     fn dtype() {

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

@@ -7,11 +7,11 @@ use std::sync::Arc;
 
 use parking_lot::Mutex;
 use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 use vortex_scalar::{Scalar, ScalarValue};
 
 use crate::arrays::ConstantArray;
-use crate::compute::{Operator, compare};
+use crate::compute::{compare, Operator};
 use crate::expr::traversal::{NodeExt, NodeVisitor, TraversalOrder};
 use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
 use crate::{Array, ArrayRef, IntoArray};
@@ -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.rs

@@ -6,7 +6,7 @@ use std::ops::Not;
 
 use prost::Message;
 use vortex_dtype::{DType, FieldName, FieldPath, Nullability};
-use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use vortex_proto::expr as pb;
 
 use crate::compute::mask;
@@ -97,23 +97,23 @@ impl VTable for GetItem {
     fn stat_max(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        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,
+        catalog: &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,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
         catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), Stat::NaNCount)
     }

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

@@ -5,7 +5,7 @@ use std::fmt::Formatter;
 use std::ops::Not;
 
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{vortex_bail, VortexResult};
 use vortex_mask::Mask;
 
 use crate::arrays::{BoolArray, ConstantArray};
@@ -72,8 +72,9 @@ impl VTable for IsNull {
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
-        catalog: &mut dyn StatsCatalog,
+        catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
+        expr.child(0).stat_nan_count(catalog);
         let field_path = expr.children()[0].stat_field_path()?;
         let null_count_expr = catalog.stats_ref(&field_path, Stat::NullCount)?;
         Some(eq(null_count_expr, lit(0u64)))
@@ -102,7 +103,6 @@ mod tests {
     use vortex_utils::aliases::hash_set::HashSet;
 
     use super::is_null;
-    use crate::IntoArray;
     use crate::arrays::{PrimitiveArray, StructArray};
     use crate::expr::exprs::binary::eq;
     use crate::expr::exprs::get_item::{col, get_item};
@@ -111,6 +111,7 @@ mod tests {
     use crate::expr::pruning::checked_pruning_expr;
     use crate::expr::test_harness;
     use crate::stats::Stat;
+    use crate::IntoArray;
 
     #[test]
     fn dtype() {

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

@@ -4,13 +4,13 @@
 use std::fmt::Formatter;
 
 use vortex_dtype::DType;
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{vortex_bail, VortexResult};
 
-use crate::ArrayRef;
 use crate::compute::list_contains as compute_list_contains;
 use crate::expr::exprs::binary::{and, gt, lt, or};
-use crate::expr::exprs::literal::{Literal, lit};
+use crate::expr::exprs::literal::{lit, Literal};
 use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::ArrayRef;
 
 pub struct ListContains;
 
@@ -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)