Merge stat lowering and row count rewriting into single pass over expressions

vortex-array/src/scalar_fn/internal/row_count.rs

@@ -5,9 +5,6 @@ use std::fmt::Formatter;
 
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
-use vortex_array::arrays::ScalarFn;
-use vortex_array::arrays::scalar_fn::ExactScalarFn;
-use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
@@ -20,28 +17,25 @@ use vortex_array::scalar_fn::ScalarFnId;
 use vortex_array::scalar_fn::ScalarFnVTable;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
-use vortex_error::vortex_ensure;
 use vortex_session::registry::CachedId;
 
 /// Zero-argument placeholder for the row count of the current evaluation scope.
 ///
-/// This is a legacy pruning hack for readers that only have a `null_count`
-/// stat and need to support `is_not_null` pruning. It is currently substituted
-/// by the zoned/file stats pruning paths before execution. New stats rewrites
-/// should prefer boolean `all_null` and `all_non_null` aggregates instead of
-/// depending on this scope-level placeholder.
+/// Stats rewrite rules emit `RowCount` when a proof needs a scope-level value that is not stored
+/// as a regular stats column — `is_not_null` is falsified by `null_count == row_count`, for
+/// example. Keeping it as a placeholder lets a rewrite rule name the row count without knowing
+/// anything about where the stats it sits beside are stored.
 ///
-/// This expression *MUST* be replaced with a concrete array before evaluation.
-/// Currently, the rewrite only happens in the context of stats pruning.
+/// It is resolved during stat binding, by [`bind_stats`], which asks the [`StatBinder`] for the
+/// row count of its scope. Binding is a single top-down pass that recurses into the expressions
+/// it substitutes, so a binder may itself emit `RowCount` and have it resolved in the same pass.
 ///
-/// `RowCount` is emitted while building pruning predicates that need a
-/// scope-level value which is not stored as a regular stats column, such as the
-/// row count of the current file or zone. The layer that owns that scope must
-/// replace each placeholder with a concrete array via [`substitute_row_count`]
-/// before evaluation.
+/// This expression *MUST* be replaced before evaluation; calling
+/// [`ScalarFnVTable::execute`] directly returns an error because this node is only a marker in a
+/// lazy expression tree.
 ///
-/// Calling [`ScalarFnVTable::execute`] directly returns an error because this
-/// node is only a marker in a lazy expression tree.
+/// [`bind_stats`]: crate::stats::bind::bind_stats
+/// [`StatBinder`]: crate::stats::bind::StatBinder
 #[derive(Clone)]
 pub struct RowCount;
 
@@ -92,66 +86,6 @@ impl ScalarFnVTable for RowCount {
     }
 }
 
-/// Returns whether `array` contains a [`RowCount`] placeholder.
-///
-/// Traversal is limited to lazy [`ScalarFnArray`] nodes produced by
-/// [`ArrayRef::apply`][crate::ArrayRef::apply]. Other arrays are evaluation
-/// leaves and cannot contain unevaluated placeholders.
-///
-/// [`ScalarFnArray`]: vortex_array::arrays::ScalarFnArray
-pub fn contains_row_count(array: &ArrayRef) -> bool {
-    if array.is::<ExactScalarFn<RowCount>>() {
-        return true;
-    }
-    match array.as_opt::<ScalarFn>() {
-        Some(view) => view.iter_children().any(contains_row_count),
-        None => false,
-    }
-}
-
-/// Replaces every [`RowCount`] placeholder with `replacement`.
-///
-/// The replacement must have the same dtype and length as each placeholder.
-/// Lazy [`ScalarFnArray`] ancestors are rewritten through slot take/put so
-/// unaffected children are preserved, while non-[`ScalarFn`] arrays are returned
-/// unchanged.
-///
-/// [`ScalarFnArray`]: vortex_array::arrays::ScalarFnArray
-pub fn substitute_row_count(array: ArrayRef, replacement: &ArrayRef) -> VortexResult<ArrayRef> {
-    if array.is::<ExactScalarFn<RowCount>>() {
-        vortex_ensure!(
-            replacement.len() == array.len(),
-            "RowCount replacement length {} does not match scope length {}",
-            replacement.len(),
-            array.len(),
-        );
-        vortex_ensure!(
-            replacement.dtype() == array.dtype(),
-            "RowCount replacement dtype {} does not match scope dtype {}",
-            replacement.dtype(),
-            array.dtype(),
-        );
-        return Ok(replacement.clone());
-    }
-
-    if !array.is::<ScalarFn>() {
-        return Ok(array);
-    }
-
-    let nchildren = array.nchildren();
-    let mut array = array;
-    for slot_idx in 0..nchildren {
-        // SAFETY: `substitute_row_count` always returns an array with the same dtype and
-        // length as its input — `RowCount` placeholders are replaced with a checked
-        // replacement (same dtype and length), and `ScalarFn` recursion preserves both by
-        // operating on each slot in place.
-        let (taken, child) = unsafe { array.take_slot_unchecked(slot_idx)? };
-        let new_child = substitute_row_count(child, replacement)?;
-        array = unsafe { taken.put_slot_unchecked(slot_idx, new_child)? };
-    }
-    Ok(array)
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_array::dtype::DType;

vortex-array/src/stats/bind.rs

@@ -12,6 +12,16 @@
 //! by a caller: zone-map field references, file-level stat literals, or typed nulls for missing
 //! stats. This lets all callers share the same falsification rules while keeping layout-specific
 //! stat storage behind [`StatBinder`].
+//!
+//! Binding also resolves [`RowCount`] placeholders, which rewrite rules emit when a proof needs
+//! the number of rows the scope covers rather than a stored statistic.
+//!
+//! Rewrite rules are independent and are combined with `or`, so several of them may prove the same
+//! thing through different statistics — `is_not_null` is falsified both by
+//! `null_count == row_count` and by `all_null`. Which of those a stats source can actually answer
+//! is only known here, and a source that answers both from the same column lowers them to the same
+//! expression. Binding collapses those duplicates on the way back up, so the predicate is not
+//! evaluated twice per row.
 
 use vortex_error::VortexResult;
 
@@ -22,7 +32,10 @@ use crate::expr::bound::lit;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::Transformed;
 use crate::scalar::Scalar;
+use crate::scalar_fn::fns::binary::Binary;
+use crate::scalar_fn::fns::operators::Operator;
 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.
 ///
@@ -42,6 +55,16 @@ pub trait StatBinder {
         stat_dtype: &DType,
     ) -> VortexResult<Option<BoundExpression>>;
 
+    /// Bind the number of rows covered by each row of the stats scope.
+    ///
+    /// This resolves the [`RowCount`] placeholders that rewrite rules emit. It is an expression
+    /// rather than a scalar because a scope may cover a different number of rows per row of its
+    /// stats table: a zone map's final zone is often shorter than the rest.
+    ///
+    /// Implementations return `Ok(None)` when the row count is unknown, and must not return an
+    /// expression that itself contains a [`RowCount`].
+    fn bind_row_count(&self) -> VortexResult<Option<BoundExpression>>;
+
     /// Expression to use when a stat is unavailable.
     ///
     /// The default is a nullable null literal, which preserves three-valued
@@ -61,19 +84,50 @@ pub fn bind_stats<B: StatBinder + ?Sized>(
     binder: &B,
 ) -> VortexResult<BoundExpression> {
     Ok(predicate
-        .transform_down(|expr| {
-            if !expr.is::<StatFn>() {
-                return Ok(Transformed::no(expr));
-            }
-
-            match bind_stat_fn(&expr, binder)? {
-                Some(bound) => Ok(Transformed::yes(bound)),
-                None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)),
-            }
-        })?
+        .transform(bind_placeholder(binder), collapse_duplicate_operand)?
         .into_inner())
 }
 
+/// Substitute a `vortex.stat` or `vortex.row_count` placeholder with the binder's representation.
+fn bind_placeholder<B: StatBinder + ?Sized>(
+    binder: &B,
+) -> impl FnMut(BoundExpression) -> VortexResult<Transformed<BoundExpression>> + '_ {
+    move |expr| {
+        // The traversal recurses into whatever it substitutes, so a binder may answer with an
+        // expression that itself contains placeholders and have them resolved in this same pass.
+        // That is what lets a stats source express `all_null` as `null_count == row_count`
+        // without a second traversal.
+        let bound = if expr.is::<StatFn>() {
+            bind_stat_fn(&expr, binder)?
+        } else if expr.is::<RowCount>() {
+            binder.bind_row_count()?
+        } else {
+            return Ok(Transformed::no(expr));
+        };
+
+        match bound {
+            Some(bound) => Ok(Transformed::yes(bound)),
+            None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)),
+        }
+    }
+}
+
+/// Collapse `a or a` and `a and a` to `a`.
+///
+/// Both are idempotent under the three-valued logic pruning uses — `null or null` is `null`, just
+/// as `null` is — so this only removes work, never changes the proof.
+fn collapse_duplicate_operand(expr: BoundExpression) -> VortexResult<Transformed<BoundExpression>> {
+    let is_duplicate = expr
+        .as_opt::<Binary>()
+        .is_some_and(|operator| matches!(operator, Operator::Or | Operator::And))
+        && expr.child(0) == expr.child(1);
+
+    if is_duplicate {
+        return Ok(Transformed::yes(expr.child(0).clone()));
+    }
+    Ok(Transformed::no(expr))
+}
+
 fn bind_stat_fn(
     expr: &BoundExpression,
     binder: &(impl StatBinder + ?Sized),
@@ -95,17 +149,23 @@ mod tests {
     use vortex_error::VortexResult;
 
     use super::*;
+    use crate::aggregate_fn::fns::all_nan::AllNan;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::dtype::StructFields;
     use crate::expr::and;
     use crate::expr::col;
+    use crate::expr::eq;
     use crate::expr::get_item;
     use crate::expr::is_null;
     use crate::expr::lit;
     use crate::expr::or;
     use crate::expr::root;
     use crate::expr::stats::Stat;
+    use crate::scalar_fn::EmptyOptions;
+    use crate::scalar_fn::ScalarFnVTableExt;
+    use crate::scalar_fn::internal::row_count::RowCount as RowCountFn;
+    use crate::stats::all_nan;
     use crate::stats::all_non_nan;
     use crate::stats::nan_count;
 
@@ -144,6 +204,19 @@ mod tests {
             aggregate_fn: &AggregateFnRef,
             _stat_dtype: &DType,
         ) -> VortexResult<Option<BoundExpression>> {
+            // `all_nan` is not stored, but it is `nan_count == row_count`. Answering with an
+            // expression that still contains a `RowCount` exercises the binding pass recursing
+            // into what it substitutes.
+            if aggregate_fn.is::<AllNan>() && self.bind_nan_count {
+                return Ok(Some(
+                    eq(
+                        get_item("f_nan_count", root()),
+                        RowCountFn.new_expr(EmptyOptions, []),
+                    )
+                    .bind(&self.stats_scope)?,
+                ));
+            }
+
             let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else {
                 return Ok(None);
             };
@@ -156,6 +229,10 @@ mod tests {
                 Ok(None)
             }
         }
+
+        fn bind_row_count(&self) -> VortexResult<Option<BoundExpression>> {
+            lit(10u64).bind(&self.stats_scope).map(Some)
+        }
     }
 
     #[test]
@@ -181,6 +258,58 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn binder_emitted_row_count_resolves_in_the_same_pass() -> VortexResult<()> {
+        let binder = TestBinder::new(true);
+
+        let bound = bind_stats(all_nan(col("f")).bind(&binder.input_scope)?, &binder)?;
+
+        assert_eq!(
+            bound,
+            eq(col("f_nan_count"), lit(10u64)).bind(&binder.stats_scope)?
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn duplicate_proofs_collapse() -> VortexResult<()> {
+        let binder = TestBinder::new(true);
+
+        // Two independent proofs of the same fact, reached through different placeholders: one
+        // states `nan_count == row_count` directly, the other asks for `all_nan`, which this
+        // binder answers the same way.
+        let predicate = or(
+            eq(nan_count(col("f")), RowCountFn.new_expr(EmptyOptions, [])),
+            all_nan(col("f")),
+        );
+        let bound = bind_stats(predicate.bind(&binder.input_scope)?, &binder)?;
+
+        assert_eq!(
+            bound,
+            eq(col("f_nan_count"), lit(10u64)).bind(&binder.stats_scope)?
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn distinct_proofs_are_both_kept() -> VortexResult<()> {
+        let binder = TestBinder::new(true);
+
+        // Only collapse operands that are actually equal.
+        let predicate = or(eq(nan_count(col("f")), lit(0u64)), all_nan(col("f")));
+        let bound = bind_stats(predicate.bind(&binder.input_scope)?, &binder)?;
+
+        assert_eq!(
+            bound,
+            or(
+                eq(col("f_nan_count"), lit(0u64)),
+                eq(col("f_nan_count"), lit(10u64)),
+            )
+            .bind(&binder.stats_scope)?
+        );
+        Ok(())
+    }
+
     #[test]
     fn missing_stats_bind_to_null_without_reducing() -> VortexResult<()> {
         let binder = TestBinder::new(false);