Bind abstract aggregate stats to legacy counts

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

Author: gatesn

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

@@ -229,7 +229,6 @@ mod tests {
     use crate::expr::pruning::field_path_stat_field_name;
     use crate::expr::root;
     use crate::expr::stats::Stat;
-    use crate::scalar::Scalar;
     use crate::scalar_fn::fns::between::BetweenOptions;
     use crate::scalar_fn::fns::between::StrictComparison;
     use crate::stats::session::StatsSession;
@@ -591,7 +590,7 @@ mod tests {
             &HashMap::from_iter([
                 (
                     FieldPath::from_name("float_col"),
-                    HashSet::from_iter([Stat::Max])
+                    HashSet::from_iter([Stat::Max, Stat::NaNCount])
                 ),
                 (
                     FieldPath::from_name("int_col"),
@@ -603,7 +602,7 @@ mod tests {
             &converted,
             &or(
                 and(
-                    lit(Scalar::null(DType::Bool(Nullability::Nullable))),
+                    eq(col("float_col_nan_count"), lit(0u64)),
                     lt_eq(col("float_col_max"), lit(10f32)),
                 ),
                 gt_eq(col("int_col_min"), lit(10)),

vortex-array/src/scalar_fn/fns/is_not_null.rs

@@ -118,7 +118,6 @@ mod tests {
     use crate::expr::eq;
     use crate::expr::get_item;
     use crate::expr::is_not_null;
-    use crate::expr::lit;
     use crate::expr::or;
     use crate::expr::pruning::checked_pruning_expr;
     use crate::expr::root;
@@ -266,7 +265,7 @@ mod tests {
             &pruning_expr,
             &or(
                 eq(col("a_null_count"), RowCount.new_expr(EmptyOptions, [])),
-                lit(Scalar::null(DType::Bool(Nullability::Nullable))),
+                eq(col("a_null_count"), RowCount.new_expr(EmptyOptions, [])),
             )
         );
         assert_eq!(

vortex-array/src/scalar_fn/fns/is_null.rs

@@ -251,7 +251,7 @@ mod tests {
             &pruning_expr,
             &or(
                 eq(col("a_null_count"), lit(0u64)),
-                lit(Scalar::null(DType::Bool(Nullability::Nullable))),
+                eq(col("a_null_count"), lit(0u64)),
             )
         );
         assert_eq!(

vortex-array/src/stats/bind.rs

@@ -11,14 +11,22 @@
 use vortex_error::VortexResult;
 
 use crate::aggregate_fn::AggregateFnRef;
+use crate::aggregate_fn::fns::all_nan::AllNan;
+use crate::aggregate_fn::fns::all_non_nan::AllNonNan;
+use crate::aggregate_fn::fns::all_non_null::AllNonNull;
+use crate::aggregate_fn::fns::all_null::AllNull;
 use crate::dtype::DType;
 use crate::expr::Expression;
+use crate::expr::eq;
 use crate::expr::lit;
 use crate::expr::stats::Stat;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::Transformed;
 use crate::scalar::Scalar;
+use crate::scalar_fn::EmptyOptions;
+use crate::scalar_fn::ScalarFnVTableExt;
 use crate::scalar_fn::fns::stat::StatFn;
+use crate::scalar_fn::internal::row_count::RowCount;
 
 /// A target that can bind abstract statistics to concrete expressions.
 pub trait StatBinder {
@@ -56,12 +64,53 @@ pub trait StatBinder {
         aggregate_fn: &AggregateFnRef,
         stat_dtype: &DType,
     ) -> VortexResult<Option<Expression>> {
+        if aggregate_fn.is::<AllNan>() {
+            let Some(nan_count) = self.bind_legacy_stat(input, Stat::NaNCount)? else {
+                return Ok(None);
+            };
+            return Ok(Some(eq(nan_count, RowCount.new_expr(EmptyOptions, []))));
+        }
+
+        if aggregate_fn.is::<AllNonNan>() {
+            let Some(nan_count) = self.bind_legacy_stat(input, Stat::NaNCount)? else {
+                return Ok(None);
+            };
+            return Ok(Some(eq(nan_count, lit(0u64))));
+        }
+
+        if aggregate_fn.is::<AllNull>() {
+            let Some(null_count) = self.bind_legacy_stat(input, Stat::NullCount)? else {
+                return Ok(None);
+            };
+            return Ok(Some(eq(null_count, RowCount.new_expr(EmptyOptions, []))));
+        }
+
+        if aggregate_fn.is::<AllNonNull>() {
+            let Some(null_count) = self.bind_legacy_stat(input, Stat::NullCount)? else {
+                return Ok(None);
+            };
+            return Ok(Some(eq(null_count, lit(0u64))));
+        }
+
         let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else {
             return Ok(None);
         };
         self.bind_stat(input, stat, stat_dtype)
     }
 
+    /// Bind one of the legacy stat slots for `input`.
+    fn bind_legacy_stat(
+        &mut 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);
+        };
+        self.bind_stat(input, stat, &stat_dtype)
+    }
+
     /// Expression to use when a stat is unavailable.
     ///
     /// The default is a nullable null literal, which preserves three-valued

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

@@ -346,7 +346,7 @@ mod tests {
     }
 
     #[test]
-    fn all_null_stat_fn_lowers_to_unknown_mask() {
+    fn all_null_stat_fn_uses_null_count() {
         let zone_map = ZoneMap::try_new(
             PType::U64.into(),
             StructArray::from_fields(&[(
@@ -363,12 +363,12 @@ mod tests {
         let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
-            BoolArray::from_iter([false, false, false])
+            BoolArray::from_iter([false, true, true])
         );
     }
 
     #[test]
-    fn all_non_null_stat_fn_lowers_to_unknown_mask() {
+    fn all_non_null_stat_fn_uses_null_count() {
         let zone_map = ZoneMap::try_new(
             PType::U64.into(),
             StructArray::from_fields(&[(
@@ -385,7 +385,7 @@ mod tests {
         let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
-            BoolArray::from_iter([false, false, false])
+            BoolArray::from_iter([true, false, false])
         );
     }
 
@@ -485,7 +485,7 @@ mod tests {
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
-            BoolArray::from_iter([false, false, false])
+            BoolArray::from_iter([true, false, false])
         );
     }
 
@@ -527,7 +527,7 @@ mod tests {
         let pruning_expr = falsify(&expr, PType::F32.into());
 
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
-        assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, false]));
+        assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true]));
     }
 
     #[test]