Push ScalarFn after Filter/Slice(x) if ScalarFn(x) has a registered kernel

vortex-array/src/arrays/filter/kernel.rs

@@ -8,6 +8,9 @@
 //! [`FilterExecuteAdaptor`] bridge these into the execution model as
 //! [`ArrayParentReduceRule`] and [`ExecuteParentKernel`] respectively.
 
+use std::sync::Arc;
+
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_session::VortexSession;
@@ -19,17 +22,33 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::array::VTable;
+use crate::arrays::Constant;
+use crate::arrays::ConstantArray;
 use crate::arrays::Dict;
 use crate::arrays::Filter;
+use crate::arrays::FilterArray;
+use crate::arrays::ScalarFn;
+use crate::arrays::ScalarFnArray;
 use crate::arrays::dict::TakeExecuteAdaptor;
+use crate::arrays::filter::FilterSlots;
+use crate::arrays::scalar_fn::ScalarFnArrayExt;
+use crate::execute_parent_for_child;
 use crate::kernel::ExecuteParentKernel;
 use crate::matcher::Matcher;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::scalar_fn::ScalarFnPlugin;
+use crate::scalar_fn::fns::between::Between;
+use crate::scalar_fn::fns::binary::Binary;
+use crate::scalar_fn::fns::fill_null::FillNull;
 
 pub(crate) fn initialize(session: &VortexSession) {
     let kernels = session.kernels();
     kernels.register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter));
+
+    for parent in [Binary.id(), Between.id(), FillNull.id()] {
+        kernels.register_execute_parent_kernel(parent, Filter, FilterScalarFnUnaryPushDownRule);
+    }
 }
 
 pub trait FilterReduce: VTable {
@@ -135,3 +154,80 @@ where
         <V as FilterKernel>::filter(array, parent.filter_mask(), ctx)
     }
 }
+
+/// If we have one non-constant child, and ScalarFn has a registered
+/// kernel for the encoding of "inner" (e.g. can operate on compressed data),
+/// it's faster to execute Fn on "inner" directly so that the we avoid
+/// canonicalizing data while executing "inner"
+///
+/// Returns Fn(inner, consts) with Fn applied eagerly.
+///
+/// Example: clickbench q1, SELECT * FROM hits WHERE AdvEngineID <> 0;
+/// AdvEngineID <> 0 is Binary over Sparse, but FlatReader applies
+/// Filter over Sparse, so we get Binary(Filter(Sparse)). Filter(Sparse)
+/// canonicalizes.
+fn scalar_fn_on_inner(
+    inner: &ArrayRef,
+    parent: ArrayView<'_, ScalarFn>,
+    child_idx: usize,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Option<ArrayRef>> {
+    if parent
+        .iter_children()
+        .filter(|c| !c.is::<Constant>())
+        .count()
+        != 1
+    {
+        return Ok(None);
+    }
+
+    let inner_len = inner.len();
+    // (inner, consts)
+    let new_children: Vec<_> = parent
+        .iter_children()
+        .map(|c| match c.as_constant() {
+            Some(scalar) => ConstantArray::new(scalar, inner_len).into_array(),
+            // by above check this is the only non-const argument
+            None => inner.clone(),
+        })
+        .collect();
+
+    // Fn(inner, consts)
+    let new_fn = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?.into_array();
+
+    // Eagerly execute swapped Fn to avoid infinite runtime with
+    // ScalarFnUnaryFilterPushDownRule. Returns Fn(inner, consts)
+    let kernels = Arc::clone(&ctx.execute_parent_kernels);
+    execute_parent_for_child(
+        "filter_scalar_fn_pushdown",
+        &new_fn, // parent, Fn
+        inner,   // child
+        child_idx,
+        &kernels,
+        ctx,
+    )
+}
+
+#[derive(Debug)]
+struct FilterScalarFnUnaryPushDownRule;
+
+impl ExecuteParentKernel<Filter> for FilterScalarFnUnaryPushDownRule {
+    type Parent = ScalarFn;
+    fn execute_parent(
+        &self,
+        child: ArrayView<'_, Filter>,
+        parent: ArrayView<'_, ScalarFn>,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let inner = child.slots()[FilterSlots::CHILD]
+            .as_ref()
+            .vortex_expect("no child for Filter");
+        let Some(new_child) = scalar_fn_on_inner(inner, parent, child_idx, ctx)? else {
+            return Ok(None);
+        };
+        let mask = child.filter_mask().clone();
+        let new_parent = FilterArray::try_new(new_child, mask)?;
+        Ok(Some(new_parent.into_array()))
+    }
+}

vortex-array/src/arrays/slice/mod.rs

@@ -14,6 +14,7 @@ mod slice_;
 mod vtable;
 
 use std::ops::Range;
+use std::sync::Arc;
 
 pub use array::SliceArraySlotsExt;
 pub use array::SliceData;
@@ -29,6 +30,7 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::array::VTable;
+use crate::execute_parent_for_child;
 use crate::kernel::ExecuteParentKernel;
 use crate::matcher::Matcher;
 use crate::optimizer::rules::ArrayParentReduceRule;
@@ -104,6 +106,31 @@ where
     }
 }
 
+/// Slice an array and eagerly resolve slice through encoding's kernels
+pub fn slice_execute(
+    array: &ArrayRef,
+    range: Range<usize>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    let sliced = array.slice(range)?;
+    let Some(view) = sliced.as_opt::<Slice>() else {
+        return Ok(sliced);
+    };
+    let Some(inner) = view.slots()[SliceSlots::CHILD].as_ref().cloned() else {
+        return Ok(sliced);
+    };
+    let kernels = Arc::clone(&ctx.execute_parent_kernels);
+    Ok(execute_parent_for_child(
+        "eager_slice",
+        &sliced,
+        &inner,
+        SliceSlots::CHILD,
+        &kernels,
+        ctx,
+    )?
+    .unwrap_or(sliced))
+}
+
 /// Adaptor that wraps a [`SliceKernel`] impl as an [`ExecuteParentKernel`].
 #[derive(Default, Debug)]
 pub struct SliceExecuteAdaptor<V>(pub V);

vortex-array/src/executor.rs

@@ -353,7 +353,7 @@ pub struct ExecutionCtx {
     session: VortexSession,
     // OnceLock avoids cloning the session allocator when a context does not allocate.
     allocator: OnceLock<BufferAllocatorRef>,
-    execute_parent_kernels: Arc<ParentExecutionKernels>,
+    pub(crate) execute_parent_kernels: Arc<ParentExecutionKernels>,
     #[cfg(debug_assertions)]
     id: usize,
     #[cfg(debug_assertions)]
@@ -621,7 +621,7 @@ fn finalize_done(
     Ok((output, None))
 }
 
-fn execute_parent_for_child(
+pub(crate) fn execute_parent_for_child(
     _phase: &'static str,
     parent: &ArrayRef,
     child: &ArrayRef,

vortex-layout/src/layouts/flat/reader.rs

@@ -11,6 +11,7 @@ use tracing::trace;
 use vortex_array::ArrayRef;
 use vortex_array::MaskFuture;
 use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::slice::slice_execute;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldMask;
 use vortex_array::expr::BoundExpression;
@@ -142,10 +143,11 @@ impl LayoutReader for FlatReader {
             //  to evaluating the expression.
             let mut array = array.clone().await?;
             let mask = mask.await?;
+            let mut ctx = session.create_execution_ctx();
 
             // Slice the array based on the row mask.
             if row_range.start > 0 || row_range.end < array.len() {
-                array = array.slice(row_range.clone())?;
+                array = slice_executed(&array, row_range.clone(), &mut ctx)?;
             }
 
             let mask_density = mask.density();
@@ -155,14 +157,12 @@ impl LayoutReader for FlatReader {
                 // after this.
                 let array = array.apply_bound(&expr)?;
                 let array = array.filter(mask.clone())?;
-                let mut ctx = session.create_execution_ctx();
                 let array_mask = array.null_as_false().execute(&mut ctx)?;
 
                 mask.intersect_by_rank(&array_mask)
             } else {
                 // Run over the full array, with a simpler bitand at the end.
                 let array = array.apply_bound(&expr)?;
-                let mut ctx = session.create_execution_ctx();
                 let array_mask = array.null_as_false().execute(&mut ctx)?;
 
                 mask.bitand(&array_mask)
@@ -193,6 +193,7 @@ impl LayoutReader for FlatReader {
         let name = Arc::clone(&self.name);
         let array = self.array_future();
         let expr = expr.clone();
+        let session = self.session.clone();
 
         Ok(async move {
             trace!("Flat array evaluation {} - {}", name, expr);
@@ -202,7 +203,8 @@ impl LayoutReader for FlatReader {
 
             // Slice the array based on the row mask.
             if row_range.start > 0 || row_range.end < array.len() {
-                array = array.slice(row_range.clone())?;
+                let mut ctx = session.create_execution_ctx();
+                array = slice_executed(&array, row_range.clone(), &mut ctx)?;
             }
 
             // First apply the filter to the array.