Stat expression

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

Author: gatesn

vortex-array/src/expr/expression.rs

@@ -10,10 +10,10 @@ use std::sync::Arc;
 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.
 ///
@@ -168,10 +168,27 @@ impl Expression {
     /// 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`.
+    ///
+    /// 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)
     }
 
+    /// Returns an expression representing the zoned maximum statistic, if available.
+    ///
+    /// See [`Self::stat_expression`] for details.
+    pub fn stat_min(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
+        self.stat_expression(Stat::Min, catalog)
+    }
+
+    /// Returns an expression representing the zoned maximum statistic, if available.
+    ///
+    /// 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.
     ///
     /// Since this is a recursive formatter, it is exposed on the public Expression type.

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::{vortex_bail, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_proto::expr as pb;
 
-use crate::compute::{between as between_compute, BetweenOptions};
+use crate::ArrayRef;
+use crate::compute::{BetweenOptions, between as between_compute};
 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.
 ///

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

@@ -5,15 +5,16 @@ use std::fmt::Formatter;
 
 use prost::Message;
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 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::{compute, ArrayRef};
+use crate::stats::Stat;
+use crate::{ArrayRef, compute};
 
 pub struct Binary;
 
@@ -127,9 +128,9 @@ impl VTable for Binary {
             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);
 
@@ -510,7 +511,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::{test_harness, Expression};
+    use crate::expr::{Expression, test_harness};
 
     #[test]
     fn and_collect_left_assoc() {

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_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
+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::ArrayRef;
+use crate::stats::Stat;
 
 /// A cast expression that converts values to a target data type.
 pub struct Cast;
@@ -86,32 +87,32 @@ impl VTable for Cast {
         })
     }
 
-    fn stat_max(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &dyn StatsCatalog,
-    ) -> Option<Expression> {
-        expr.children()[0].stat_max(catalog)
-    }
-
-    fn stat_min(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &dyn StatsCatalog,
-    ) -> Option<Expression> {
-        expr.children()[0].stat_min(catalog)
-    }
-
-    fn stat_nan_count(
+    fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
+        stat: Stat,
         catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        expr.children()[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::Max
+            | Stat::Min
+            | Stat::NaNCount
+            | Stat::Sum
+            | Stat::UncompressedSizeInBytes => expr.child(0).stat_expression(stat, catalog),
+            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
+            }
+        }
     }
 }
 
@@ -136,11 +137,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::{test_harness, Expression};
-    use crate::IntoArray;
+    use crate::expr::{Expression, test_harness};
 
     #[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::{vortex_bail, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_scalar::{Scalar, ScalarValue};
 
 use crate::arrays::ConstantArray;
-use crate::compute::{compare, Operator};
+use crate::compute::{Operator, compare};
 use crate::expr::traversal::{NodeExt, NodeVisitor, TraversalOrder};
 use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
 use crate::{Array, ArrayRef, IntoArray};

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::{vortex_bail, vortex_err, VortexResult};
+use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr as pb;
 
 use crate::compute::mask;
@@ -94,34 +94,20 @@ impl VTable for GetItem {
         }
     }
 
-    fn stat_max(
+    fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
+        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: &dyn StatsCatalog,
-    ) -> Option<Expression> {
-        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), Stat::Min)
-    }
-
-    fn stat_nan_count(
-        &self,
-        expr: &ExpressionView<Self>,
-        catalog: &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): For now, we maintain the existing *broken* behavior of assuming that
+        //  getitem refers to a field on the root scope.
+        catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), stat)
     }
 }
 

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::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_mask::Mask;
 
 use crate::arrays::{BoolArray, ConstantArray};
@@ -74,9 +74,7 @@ impl VTable for IsNull {
         expr: &ExpressionView<Self>,
         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)?;
+        let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?;
         Some(eq(null_count_expr, lit(0u64)))
     }
 }
@@ -103,6 +101,7 @@ 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,7 +110,6 @@ 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::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_bail};
 
+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::{lit, Literal};
+use crate::expr::exprs::literal::{Literal, lit};
 use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
-use crate::ArrayRef;
 
 pub struct ListContains;