diff --git a/vortex-array/src/arrays/filter/kernel.rs b/vortex-array/src/arrays/filter/kernel.rs index a62dc7910a0..02c6ef43634 100644 --- a/vortex-array/src/arrays/filter/kernel.rs +++ b/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 ::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> { + if parent + .iter_children() + .filter(|c| !c.is::()) + .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 for FilterScalarFnUnaryPushDownRule { + type Parent = ScalarFn; + fn execute_parent( + &self, + child: ArrayView<'_, Filter>, + parent: ArrayView<'_, ScalarFn>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + 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())) + } +} diff --git a/vortex-array/src/arrays/slice/mod.rs b/vortex-array/src/arrays/slice/mod.rs index bdbd75f45b6..699cfe1fde9 100644 --- a/vortex-array/src/arrays/slice/mod.rs +++ b/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, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let sliced = array.slice(range)?; + let Some(view) = sliced.as_opt::() 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(pub V); diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index dab9ccfcfb5..573716239ef 100644 --- a/vortex-array/src/executor.rs +++ b/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, - execute_parent_kernels: Arc, + pub(crate) execute_parent_kernels: Arc, #[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, diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index aa7609f1659..6d83afd6c2f 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/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.