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

vortex-array/src/expr/bound_expression.rs

@@ -143,7 +143,12 @@ impl BoundExpression {
         children: impl IntoIterator<Item = BoundExpression>,
     ) -> VortexResult<Self> {
         let children = Vec::from_iter(children);
-        let BoundExpression::Scalar { scalar_fn, .. } = &self else {
+        let BoundExpression::Scalar {
+            dtype,
+            scalar_fn,
+            children: old_children,
+        } = &self
+        else {
             vortex_ensure!(
                 children.is_empty(),
                 "Root expression cannot have {} children",
@@ -152,6 +157,23 @@ impl BoundExpression {
             return Ok(self);
         };
 
+        // A return dtype is a function of the argument dtypes alone, so replacing children that
+        // type the same cannot change it. Rewrites usually substitute deep inside a tree and leave
+        // every dtype on the path to the root untouched, and recomputing one means a vector of
+        // cloned dtypes and a virtual call per node on that path.
+        if children.len() == old_children.len()
+            && children
+                .iter()
+                .zip(old_children.iter())
+                .all(|(new, old)| new.dtype() == old.dtype())
+        {
+            return Ok(Self::Scalar {
+                dtype: dtype.clone(),
+                scalar_fn: scalar_fn.clone(),
+                children: children.into(),
+            });
+        }
+
         Self::try_new(scalar_fn.clone(), children)
     }
 

vortex-array/src/expr/traversal/mod.rs

@@ -556,28 +556,38 @@ impl Node for BoundExpression {
         };
 
         let mut order = TraversalOrder::Continue;
-        let mut changed = false;
-        let children = children
-            .iter()
-            .cloned()
-            .map(|child| match order {
-                TraversalOrder::Continue | TraversalOrder::Skip => f(child).map(|result| {
+        // Stays `None` until a child actually changes. Most nodes of a rewritten tree are
+        // untouched, and collecting their children into a vector only to discard it is the
+        // dominant cost of a rewrite over a large predicate.
+        let mut rewritten: Option<Vec<Self>> = None;
+
+        for (index, child) in children.iter().enumerate() {
+            let value = match order {
+                TraversalOrder::Continue | TraversalOrder::Skip => {
+                    let result = f(child.clone())?;
                     order = result.order;
-                    changed |= result.changed;
+                    if result.changed && rewritten.is_none() {
+                        let mut prefix = Vec::with_capacity(children.len());
+                        prefix.extend_from_slice(&children[..index]);
+                        rewritten = Some(prefix);
+                    }
                     result.value
-                }),
-                TraversalOrder::Stop => Ok(child),
-            })
-            .collect::<VortexResult<Vec<_>>>()?;
+                }
+                TraversalOrder::Stop => child.clone(),
+            };
 
-        if changed {
-            Ok(Transformed {
-                value: self.with_children(children)?,
+            if let Some(rewritten) = &mut rewritten {
+                rewritten.push(value);
+            }
+        }
+
+        match rewritten {
+            Some(rewritten) => Ok(Transformed {
+                value: self.with_children(rewritten)?,
                 order,
                 changed: true,
-            })
-        } else {
-            Ok(Transformed::no(self))
+            }),
+            None => Ok(Transformed::no(self)),
         }
     }
 

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;