Prunning expressions can reference rowcount (#7589)

This lets us effectively prune expressions like IsNotNull

fix #7187

---------

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

Cargo.lock

@@ -1,4 +1,4 @@
 [toolchain]
 channel = "1.91.0"
 components = ["rust-src", "rustfmt", "clippy", "rust-analyzer"]
-profile = "minimal"
+profile = "minimal"

java/testfiles/Cargo.lock

@@ -17532,6 +17532,58 @@ pub trait vortex_array::scalar_fn::fns::zip::ZipReduce: vortex_array::VTable
 
 pub fn vortex_array::scalar_fn::fns::zip::ZipReduce::zip(array: vortex_array::ArrayView<'_, Self>, if_false: &vortex_array::ArrayRef, mask: &vortex_array::ArrayRef) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
 
+pub mod vortex_array::scalar_fn::internal
+
+pub mod vortex_array::scalar_fn::internal::row_count
+
+pub struct vortex_array::scalar_fn::internal::row_count::RowCount
+
+impl core::clone::Clone for vortex_array::scalar_fn::internal::row_count::RowCount
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::clone(&self) -> vortex_array::scalar_fn::internal::row_count::RowCount
+
+impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::internal::row_count::RowCount
+
+pub type vortex_array::scalar_fn::internal::row_count::RowCount::Options = vortex_array::scalar_fn::EmptyOptions
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::child_name(&self, _options: &Self::Options, _child_idx: usize) -> vortex_array::scalar_fn::ChildName
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::fmt_sql(&self, _options: &Self::Options, _expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::id(&self) -> vortex_array::scalar_fn::ScalarFnId
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::is_fallible(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::is_null_sensitive(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::reduce(&self, options: &Self::Options, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::scalar_fn::ReduceNodeRef>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::return_dtype(&self, _options: &Self::Options, _args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::contains_row_count(array: &vortex_array::ArrayRef) -> bool
+
+pub fn vortex_array::scalar_fn::internal::row_count::substitute_row_count(array: vortex_array::ArrayRef, replacement: &vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
 pub mod vortex_array::scalar_fn::session
 
 pub struct vortex_array::scalar_fn::session::ScalarFnSession
@@ -18680,6 +18732,44 @@ pub fn vortex_array::scalar_fn::fns::zip::Zip::stat_falsification(&self, options
 
 pub fn vortex_array::scalar_fn::fns::zip::Zip::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
 
+impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::internal::row_count::RowCount
+
+pub type vortex_array::scalar_fn::internal::row_count::RowCount::Options = vortex_array::scalar_fn::EmptyOptions
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::child_name(&self, _options: &Self::Options, _child_idx: usize) -> vortex_array::scalar_fn::ChildName
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::fmt_sql(&self, _options: &Self::Options, _expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::id(&self) -> vortex_array::scalar_fn::ScalarFnId
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::is_fallible(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::is_null_sensitive(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::reduce(&self, options: &Self::Options, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::scalar_fn::ReduceNodeRef>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::return_dtype(&self, _options: &Self::Options, _args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
 pub trait vortex_array::scalar_fn::ScalarFnVTableExt: vortex_array::scalar_fn::ScalarFnVTable
 
 pub fn vortex_array::scalar_fn::ScalarFnVTableExt::bind(&self, options: Self::Options) -> vortex_array::scalar_fn::ScalarFnRef

rust-toolchain.toml

@@ -86,6 +86,11 @@ pub fn field_path_stat_field_name(field_path: &FieldPath, stat: Stat) -> FieldNa
 /// cannot hold, and false if it cannot be determined from stats alone whether the positions can
 /// be pruned.
 ///
+/// Some rewrites, such as `is_not_null(...)`, emit
+/// [`row_count`][crate::scalar_fn::internal::row_count] placeholders. The evaluation layer must
+/// replace those placeholders with the row count for its current scope before
+/// executing the returned expression.
+///
 /// If the falsification logic attempts to access an unknown stat,
 /// this function will return `None`.
 pub fn checked_pruning_expr(

vortex-array/public-api.lock

@@ -3,6 +3,7 @@
 
 use std::fmt::Formatter;
 
+use vortex_array::scalar_fn::internal::row_count::RowCount;
 use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 
@@ -14,17 +15,15 @@ use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::expr::Expression;
 use crate::expr::StatsCatalog;
-use crate::expr::and;
 use crate::expr::eq;
-use crate::expr::gt;
-use crate::expr::lit;
 use crate::expr::stats::Stat;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
 use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::ExecutionArgs;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::ScalarFnVTable;
+use crate::scalar_fn::ScalarFnVTableExt;
 use crate::validity::Validity;
 
 /// Expression that checks for non-null values.
@@ -106,40 +105,42 @@ impl ScalarFnVTable for IsNotNull {
         expr: &Expression,
         catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        // is_not_null is falsified when ALL values are null, i.e. null_count == len.
-        // Since there is no len stat in the zone map, we approximate using IsConstant:
-        // if the zone is constant and has any nulls, then all values must be null.
-        //
-        // TODO(#7187): Add a len stat to enable the more general falsification:
-        //   null_count == len => is_not_null is all false.
-        let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?;
-        let is_constant_expr = expr.child(0).stat_expression(Stat::IsConstant, catalog)?;
-        // If the zone is constant (is_constant == true) and has nulls (null_count > 0),
-        // then all values must be null, so is_not_null is all false.
-        Some(and(
-            eq(is_constant_expr, lit(true)),
-            gt(null_count_expr, lit(0u64)),
-        ))
+        // is_not_null is falsified when ALL values are null, i.e. null_count == row_count.
+        let child = expr.child(0);
+        let null_count_expr = child.stat_expression(Stat::NullCount, catalog)?;
+        Some(eq(null_count_expr, RowCount.new_expr(EmptyOptions, [])))
     }
 }
 
 #[cfg(test)]
 mod tests {
     use vortex_buffer::buffer;
     use vortex_error::VortexExpect as _;
+    use vortex_utils::aliases::hash_map::HashMap;
+    use vortex_utils::aliases::hash_set::HashSet;
 
     use crate::IntoArray;
     use crate::LEGACY_SESSION;
     use crate::VortexSessionExecute;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::StructArray;
     use crate::dtype::DType;
+    use crate::dtype::Field;
+    use crate::dtype::FieldPath;
+    use crate::dtype::FieldPathSet;
     use crate::dtype::Nullability;
+    use crate::expr::col;
+    use crate::expr::eq;
     use crate::expr::get_item;
     use crate::expr::is_not_null;
+    use crate::expr::pruning::checked_pruning_expr;
     use crate::expr::root;
+    use crate::expr::stats::Stat;
     use crate::expr::test_harness;
     use crate::scalar::Scalar;
+    use crate::scalar_fn::EmptyOptions;
+    use crate::scalar_fn::internal::row_count::RowCount;
+    use crate::scalar_fn::vtable::ScalarFnVTableExt;
 
     #[test]
     fn dtype() {
@@ -255,50 +256,29 @@ mod tests {
 
     #[test]
     fn test_is_not_null_sensitive() {
-        use crate::expr::col;
         assert!(is_not_null(col("a")).signature().is_null_sensitive());
     }
 
     #[test]
     fn test_is_not_null_falsification() {
-        use vortex_utils::aliases::hash_map::HashMap;
-        use vortex_utils::aliases::hash_set::HashSet;
-
-        use crate::dtype::Field;
-        use crate::dtype::FieldPath;
-        use crate::dtype::FieldPathSet;
-        use crate::expr::and;
-        use crate::expr::col;
-        use crate::expr::eq;
-        use crate::expr::gt;
-        use crate::expr::lit;
-        use crate::expr::pruning::checked_pruning_expr;
-        use crate::expr::stats::Stat;
-
         let expr = is_not_null(col("a"));
 
         let (pruning_expr, st) = checked_pruning_expr(
             &expr,
-            &FieldPathSet::from_iter([
-                FieldPath::from_iter([Field::Name("a".into()), Field::Name("null_count".into())]),
-                FieldPath::from_iter([Field::Name("a".into()), Field::Name("is_constant".into())]),
-            ]),
+            &FieldPathSet::from_iter([FieldPath::from_iter([
+                Field::Name("a".into()),
+                Field::Name("null_count".into()),
+            ])]),
         )
         .unwrap();
 
         assert_eq!(
             &pruning_expr,
-            &and(
-                eq(col("a_is_constant"), lit(true)),
-                gt(col("a_null_count"), lit(0u64)),
-            )
+            &eq(col("a_null_count"), RowCount.new_expr(EmptyOptions, []))
         );
         assert_eq!(
             st.map(),
-            &HashMap::from_iter([(
-                FieldPath::from_name("a"),
-                HashSet::from([Stat::NullCount, Stat::IsConstant])
-            )])
+            &HashMap::from_iter([(FieldPath::from_name("a"), HashSet::from([Stat::NullCount]))])
         );
     }
 }

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

@@ -0,0 +1,4 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+pub mod row_count;