Use chunked execute_parent for stat expressions

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

Author: gatesn

vortex-array/src/arrays/chunked/compute/kernel.rs

@@ -1,18 +1,77 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::array::ArrayView;
 use crate::arrays::Chunked;
+use crate::arrays::ChunkedArray;
+use crate::arrays::ScalarFn;
+use crate::arrays::ScalarFnArray;
+use crate::arrays::chunked::ChunkedArrayExt;
 use crate::arrays::dict::TakeExecuteAdaptor;
 use crate::arrays::filter::FilterExecuteAdaptor;
+use crate::arrays::scalar_fn::ExactScalarFn;
+use crate::arrays::scalar_fn::ScalarFnArrayExt;
+use crate::arrays::scalar_fn::ScalarFnArrayView;
 use crate::arrays::slice::SliceExecuteAdaptor;
+use crate::kernel::ExecuteParentKernel;
 use crate::kernel::ParentKernelSet;
 use crate::scalar_fn::fns::mask::MaskExecuteAdaptor;
+use crate::scalar_fn::fns::stat::StatFn;
 use crate::scalar_fn::fns::zip::ZipExecuteAdaptor;
 
 pub(crate) static PARENT_KERNELS: ParentKernelSet<Chunked> = ParentKernelSet::new(&[
+    ParentKernelSet::lift(&ChunkedStatExecuteParentKernel),
     ParentKernelSet::lift(&FilterExecuteAdaptor(Chunked)),
     ParentKernelSet::lift(&MaskExecuteAdaptor(Chunked)),
     ParentKernelSet::lift(&SliceExecuteAdaptor(Chunked)),
     ParentKernelSet::lift(&TakeExecuteAdaptor(Chunked)),
     ParentKernelSet::lift(&ZipExecuteAdaptor(Chunked)),
 ]);
+
+#[derive(Debug)]
+struct ChunkedStatExecuteParentKernel;
+
+impl ExecuteParentKernel<Chunked> for ChunkedStatExecuteParentKernel {
+    type Parent = ExactScalarFn<StatFn>;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, Chunked>,
+        parent: ScalarFnArrayView<'_, StatFn>,
+        child_idx: usize,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if child_idx != 0 {
+            return Ok(None);
+        }
+
+        tracing::trace!(
+            "stat({}) descending into ChunkedArray with {} chunks",
+            parent.options.aggregate_fn(),
+            array.nchunks()
+        );
+
+        let scalar_fn = parent
+            .as_opt::<ScalarFn>()
+            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray")
+            .scalar_fn()
+            .clone();
+        let chunks = array
+            .iter_chunks()
+            .map(|chunk| {
+                ScalarFnArray::try_new(scalar_fn.clone(), vec![chunk.clone()], chunk.len())
+                    .map(IntoArray::into_array)
+            })
+            .collect::<VortexResult<Vec<_>>>()?;
+
+        Ok(Some(
+            ChunkedArray::try_new(chunks, parent.dtype().clone())?.into_array(),
+        ))
+    }
+}

vortex-array/src/arrays/chunked/compute/rules.rs

@@ -21,6 +21,7 @@ use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
 use crate::scalar_fn::fns::cast::CastReduceAdaptor;
 use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor;
+use crate::scalar_fn::fns::stat::StatFn;
 
 pub(crate) const PARENT_RULES: ParentRuleSet<Chunked> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&CastReduceAdaptor(Chunked)),
@@ -44,6 +45,9 @@ impl ArrayParentReduceRule<Chunked> for ChunkedUnaryScalarFnPushDownRule {
         if parent.nchildren() != 1 {
             return Ok(None);
         }
+        if parent.scalar_fn().is::<StatFn>() {
+            return Ok(None);
+        }
 
         let new_chunks: Vec<_> = array
             .iter_chunks()
@@ -76,6 +80,10 @@ impl ArrayParentReduceRule<Chunked> for ChunkedConstantScalarFnPushDownRule {
         parent: ArrayView<'_, ScalarFn>,
         child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
+        if parent.scalar_fn().is::<StatFn>() {
+            return Ok(None);
+        }
+
         for (idx, child) in parent.iter_children().enumerate() {
             if idx == child_idx {
                 continue;

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

@@ -13,10 +13,7 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::aggregate_fn::AggregateFnRef;
-use crate::arrays::Chunked;
-use crate::arrays::ChunkedArray;
 use crate::arrays::ConstantArray;
-use crate::arrays::chunked::ChunkedArrayExt;
 use crate::dtype::DType;
 use crate::expr::Expression;
 use crate::expr::stats::StatsProvider;
@@ -115,26 +112,6 @@ impl ScalarFnVTable for StatFn {
     ) -> VortexResult<ArrayRef> {
         let input = args.get(0)?;
         let dtype = stat_dtype(options.aggregate_fn(), input.dtype())?;
-
-        // Recurse into each chunk so the output keeps per-chunk granularity (one constant
-        // per chunk) instead of collapsing to a single combined stat across the whole array.
-        // Per-chunk granularity is what makes the result useful for pruning: predicates can
-        // drop whole chunks at a time. Without this, we'd lose every chunk boundary as a
-        // pruning opportunity. Other encodings (zone maps, etc.) will need similar
-        // structure-preserving handling once they land.
-        if let Some(chunked) = input.as_opt::<Chunked>() {
-            tracing::trace!(
-                "stat({}) descending into ChunkedArray with {} chunks",
-                options.aggregate_fn(),
-                chunked.nchunks()
-            );
-            let chunks = chunked
-                .iter_chunks()
-                .map(|chunk| stat_array(chunk, options.aggregate_fn(), dtype.clone(), chunk.len()))
-                .collect::<VortexResult<Vec<_>>>()?;
-            return Ok(ChunkedArray::try_new(chunks, dtype)?.into_array());
-        }
-
         stat_array(&input, options.aggregate_fn(), dtype, args.row_count())
     }
 }

vortex-array/src/stats/expr.rs

@@ -24,6 +24,7 @@ mod tests {
     use vortex_error::VortexResult;
 
     use super::stat;
+    use crate::ArrayRef;
     use crate::Canonical;
     use crate::IntoArray;
     use crate::LEGACY_SESSION;
@@ -102,6 +103,9 @@ mod tests {
         .into_array();
 
         let result = chunked.apply(&stat(root(), AggregateFn::new(Sum, EmptyOptions).erased()))?;
+        assert!(result.as_opt::<Chunked>().is_none());
+
+        let result = result.execute::<ArrayRef>(&mut LEGACY_SESSION.create_execution_ctx())?;
 
         let chunked_result = result
             .as_opt::<Chunked>()