From 408386b489caea974425ff89807bc5b90d944112 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 16 Apr 2026 16:32:16 -0400 Subject: [PATCH 01/16] perf: optimize iterative exe Signed-off-by: Joe Isaacs --- .../src/arrays/chunked/vtable/canonical.rs | 16 +++++++--------- vortex-array/src/arrays/chunked/vtable/mod.rs | 15 ++++++++++++++- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 1e32f7295bc..2fc45e6544a 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -5,12 +5,12 @@ use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; +use crate::{AnyCanonical, ArrayRef}; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::Chunked; +use crate::arrays::{Chunked, ListView, Struct}; use crate::arrays::ChunkedArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -36,7 +36,7 @@ pub(super) fn _canonicalize( return Ok(Canonical::empty(array.dtype())); } if array.nchunks() == 1 { - return array.chunk(0).clone().execute::(ctx); + return Ok(Canonical::from(array.chunk(0).as_::())); } let owned_chunks: Vec = array.iter_chunks().cloned().collect(); @@ -46,7 +46,6 @@ pub(super) fn _canonicalize( &owned_chunks, Validity::copy_from_array(array.array())?, struct_dtype, - ctx, )?; Canonical::Struct(struct_array) } @@ -72,15 +71,14 @@ fn pack_struct_chunks( chunks: &[ArrayRef], validity: Validity, struct_dtype: &StructFields, - ctx: &mut ExecutionCtx, ) -> VortexResult { let len = chunks.iter().map(|chunk| chunk.len()).sum(); let mut field_arrays = Vec::new(); - let executed_chunks: Vec = chunks + let executed_chunks = chunks .iter() - .map(|c| c.clone().execute::(ctx)) - .collect::>()?; + .map(|c| c.clone().downcast::()) + .collect::>(); for (field_idx, field_dtype) in struct_dtype.fields().enumerate() { let mut field_chunks = Vec::with_capacity(chunks.len()); @@ -140,7 +138,7 @@ fn swizzle_list_chunks( let mut next_list = 0usize; for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; + let chunk_array = chunk.clone().downcast::(); // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make // the final output `ListView` also zero-copyable to `List`. let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact)?; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 5e9c4a4d817..42209aa1730 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -13,7 +13,7 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayEq; +use crate::{AnyCanonical, ArrayEq}; use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; @@ -239,6 +239,19 @@ impl VTable for Chunked { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + // Iteratively request execution of each chunk until it reaches canonical form. + // This gives the scheduler visibility into child execution and enables + // cross-step optimizations between chunk decoding steps. + for i in 0..array.nchunks() { + if !array.chunk(i).is::() { + return Ok(ExecutionResult::execute_slot::( + array, + i + CHUNKS_OFFSET, + )); + } + } + + // All chunks are now canonical — combine them. Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) } From 9e7ca16b823feb3ffdc7e9211b9137719ff9984f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 17 Apr 2026 09:39:19 -0400 Subject: [PATCH 02/16] add faster take/put methods on arrays Signed-off-by: Joe Isaacs --- vortex-array/src/array/erased.rs | 48 +++++++++++++++++++ vortex-array/src/array/mod.rs | 7 +++ .../src/arrays/chunked/vtable/canonical.rs | 7 ++- vortex-array/src/arrays/chunked/vtable/mod.rs | 3 +- vortex-array/src/executor.rs | 15 +++--- 5 files changed, 71 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 84367076314..a4ac4cfd00d 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -431,6 +431,54 @@ impl ArrayRef { self.with_slots(slots) } + /// Take a slot for executor-owned physical rewrites. This has the result that the array may + /// either be taken or cloned from the parent. + /// + /// The array can be put back with [`put_slot_unchecked`]. + /// + /// # Safety + /// The caller must put back a slot with the same logical dtype and length before exposing the + /// parent array, and must only use this for physical rewrites. + pub(crate) unsafe fn take_slot_unchecked( + mut self, + slot_idx: usize, + ) -> VortexResult<(ArrayRef, ArrayRef)> { + let child = if let Some(inner) = Arc::get_mut(&mut self.0) { + unsafe { inner.slots_mut()[slot_idx].take() } + .vortex_expect("take_slot_unchecked cannot take an absent slot") + } else { + self.slots()[slot_idx] + .as_ref() + .vortex_expect("take_slot_unchecked cannot take an absent slot") + .clone() + }; + + Ok((self, child)) + } + + /// Puts an array into `slot_idx` by either, cloning the inner array if the Arc is not exclusive + /// or replacing the slot in this `ArrayRef`. + /// This is the mirror of [`take_slot_unchecked`]. + /// + /// # Safety + /// The replacement must have the same logical dtype and length as the taken slot, and this + /// must only be used for physical rewrites. + pub(crate) unsafe fn put_slot_unchecked( + mut self, + slot_idx: usize, + replacement: ArrayRef, + ) -> VortexResult { + if let Some(inner) = Arc::get_mut(&mut self.0) { + unsafe { inner.slots_mut()[slot_idx] = Some(replacement) }; + return Ok(self); + } + + let mut slots = self.slots().to_vec(); + slots[slot_idx] = Some(replacement); + let inner = Arc::clone(&self.0); + inner.with_slots(self, slots) + } + /// Returns a new array with the provided slots. /// /// This is only valid for physical rewrites: slot count, presence, logical `DType`, and diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 8193c128b1e..d4ed9096f82 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -68,6 +68,9 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug { /// Returns the slots of the array. fn slots(&self) -> &[Option]; + /// Returns mutable slots of the array. + unsafe fn slots_mut(&mut self) -> &mut [Option]; + /// Returns the encoding ID of the array. fn encoding_id(&self) -> ArrayId; @@ -212,6 +215,10 @@ impl DynArray for ArrayInner { &self.slots } + unsafe fn slots_mut(&mut self) -> &mut [Option] { + &mut self.slots + } + fn encoding_id(&self) -> ArrayId { self.vtable.id() } diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 2fc45e6544a..2364e57f514 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -5,15 +5,18 @@ use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::{AnyCanonical, ArrayRef}; +use crate::AnyCanonical; +use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::{Chunked, ListView, Struct}; +use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::listview::ListViewArrayExt; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 42209aa1730..1171f40d230 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -13,7 +13,8 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::{AnyCanonical, ArrayEq}; +use crate::AnyCanonical; +use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 8e0c86ff44f..35d638487c8 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -122,7 +122,8 @@ impl ArrayRef { return Ok(current); } Some((parent, slot_idx, _)) => { - current = parent.with_slot(slot_idx, current)?.optimize()?; + current = + unsafe { parent.put_slot_unchecked(slot_idx, current)? }.optimize()?; continue; } } @@ -140,7 +141,8 @@ impl ArrayRef { )); current = rewritten.optimize()?; if let Some((parent, slot_idx, _)) = stack.pop() { - current = parent.with_slot(slot_idx, current)?.optimize()?; + current = + unsafe { parent.put_slot_unchecked(slot_idx, current)? }.optimize()?; } continue; } @@ -151,13 +153,14 @@ impl ArrayRef { match step { ExecutionStep::ExecuteSlot(i, done) => { let child = array.slots()[i] - .clone() + .as_ref() .vortex_expect("ExecuteSlot index in bounds"); ctx.log(format_args!( "ExecuteSlot({i}): pushing {}, focusing on {}", array, child )); - stack.push((array, i, done)); + let (parent, child) = unsafe { array.take_slot_unchecked(i)? }; + stack.push((parent, i, done)); current = child.optimize()?; } ExecutionStep::Done => { @@ -327,9 +330,9 @@ impl Executable for ArrayRef { ExecutionStep::ExecuteSlot(i, _) => { // For single-step execution, handle ExecuteSlot by executing the slot, // replacing it, and returning the updated array. - let child = array.slots()[i].clone().vortex_expect("valid slot index"); + let (array, child) = unsafe { array.take_slot_unchecked(i)? }; let executed_child = child.execute::(ctx)?; - array.with_slot(i, executed_child) + unsafe { array.put_slot_unchecked(i, executed_child) } } } } From 9c9cbd011a1804b85fde24950065d7ce0c8e67b8 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 10:03:22 -0400 Subject: [PATCH 03/16] add faster take/put methods on arrays Signed-off-by: Joe Isaacs --- .../src/arrays/chunked/compute/kernel.rs | 26 ++++-- .../src/arrays/chunked/compute/rules.rs | 27 ++++-- vortex-array/src/arrays/dict/compute/rules.rs | 30 ++++-- vortex-array/src/arrays/dict/vtable/kernel.rs | 22 ++++- .../src/arrays/primitive/compute/rules.rs | 21 ++++- .../src/arrays/primitive/vtable/kernel.rs | 27 ++++-- vortex-array/src/kernel.rs | 72 ++++++++++++++- vortex-array/src/optimizer/rules.rs | 91 ++++++++++++++++++- vortex-session/public-api.lock | 2 + vortex-session/src/registry.rs | 6 ++ 10 files changed, 280 insertions(+), 44 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/kernel.rs b/vortex-array/src/arrays/chunked/compute/kernel.rs index 502c1fe8337..d78f9787fb2 100644 --- a/vortex-array/src/arrays/chunked/compute/kernel.rs +++ b/vortex-array/src/arrays/chunked/compute/kernel.rs @@ -1,18 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Chunked; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&MaskExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&SliceExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Chunked)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 5] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(Chunked), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.mask"), &MaskExecuteAdaptor(Chunked)), + ParentKernelSet::lift_id(CachedId::new("vortex.slice"), &SliceExecuteAdaptor(Chunked)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Chunked)), + ParentKernelSet::lift_id(CachedId::new("vortex.zip"), &ZipExecuteAdaptor(Chunked)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/chunked/compute/rules.rs b/vortex-array/src/arrays/chunked/compute/rules.rs index 8ef63a188c3..13f0c3f0dac 100644 --- a/vortex-array/src/arrays/chunked/compute/rules.rs +++ b/vortex-array/src/arrays/chunked/compute/rules.rs @@ -3,6 +3,7 @@ use itertools::Itertools; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -18,16 +19,30 @@ use crate::arrays::scalar_fn::AnyScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(Chunked)), - ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule), - ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule), - ParentRuleSet::lift(&FillNullReduceAdaptor(Chunked)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Chunked)), + ParentRuleSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullReduceAdaptor(Chunked), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ + ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule), + ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule), + ], +); /// Push down any unary scalar function through chunked arrays. #[derive(Debug)] diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index 23fc05d08d6..c3149add437 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayRef; @@ -23,6 +24,8 @@ use crate::arrays::slice::SliceReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::optimizer::ArrayOptimizer; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastReduceAdaptor; @@ -31,15 +34,24 @@ use crate::scalar_fn::fns::mask::MaskReduceAdaptor; use crate::scalar_fn::fns::pack::Pack; use crate::validity::Validity; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&FilterReduceAdaptor(Dict)), - ParentRuleSet::lift(&CastReduceAdaptor(Dict)), - ParentRuleSet::lift(&MaskReduceAdaptor(Dict)), - ParentRuleSet::lift(&LikeReduceAdaptor(Dict)), - ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule), - ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule), - ParentRuleSet::lift(&SliceReduceAdaptor(Dict)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(Dict)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Dict)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Dict)), + ParentRuleSet::lift_id(CachedId::new("vortex.like"), &LikeReduceAdaptor(Dict)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Dict)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ + ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule), + ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule), + ], +); /// Push down a scalar function to run only over the values of a dictionary array. #[derive(Debug)] diff --git a/vortex-array/src/arrays/dict/vtable/kernel.rs b/vortex-array/src/arrays/dict/vtable/kernel.rs index fd1b0ed0a6d..6fe1cd33c9a 100644 --- a/vortex-array/src/arrays/dict/vtable/kernel.rs +++ b/vortex-array/src/arrays/dict/vtable/kernel.rs @@ -1,14 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Dict)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Dict)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Dict)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.binary"), &CompareExecuteAdaptor(Dict)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Dict)), + ParentKernelSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullExecuteAdaptor(Dict), + ), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/primitive/compute/rules.rs b/vortex-array/src/arrays/primitive/compute/rules.rs index 99b0a7464d5..11043814d32 100644 --- a/vortex-array/src/arrays/primitive/compute/rules.rs +++ b/vortex-array/src/arrays/primitive/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -11,14 +12,24 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&PrimitiveMaskedValidityRule), - ParentRuleSet::lift(&MaskReduceAdaptor(Primitive)), - ParentRuleSet::lift(&SliceReduceAdaptor(Primitive)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.masked"), &PrimitiveMaskedValidityRule), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Primitive)), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(Primitive), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Rule to push down validity masking from MaskedArray parent into PrimitiveArray child. /// diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 03b6cfa2b98..a4af02c545c 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -1,16 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Primitive; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&CastExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Primitive)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 4] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.between"), + &BetweenExecuteAdaptor(Primitive), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.cast"), &CastExecuteAdaptor(Primitive)), + ParentKernelSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullExecuteAdaptor(Primitive), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Primitive)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/kernel.rs b/vortex-array/src/kernel.rs index f5b75471437..5ad9849ddc4 100644 --- a/vortex-array/src/kernel.rs +++ b/vortex-array/src/kernel.rs @@ -16,9 +16,12 @@ use std::any::type_name; use std::fmt::Debug; use std::marker::PhantomData; +use std::sync::OnceLock; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use crate::ArrayId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::array::ArrayView; @@ -31,17 +34,64 @@ use crate::matcher::Matcher; /// a kernel whose [`Matcher`] matches the parent array type. The first matching kernel that /// returns `Some` wins. pub struct ParentKernelSet { - kernels: &'static [&'static dyn DynParentKernel], + keyed: &'static [ParentKernelEntry], + dense: Option<&'static ParentKernelDense>, + dynamic: &'static [&'static dyn DynParentKernel], } +/// A parent kernel keyed by exact parent encoding id. +pub(crate) struct ParentKernelEntry { + parent_id: CachedId, + kernel: &'static dyn DynParentKernel, +} + +pub(crate) type ParentKernelDense = OnceLock>]>>; + impl ParentKernelSet { + fn kernels_for(&self, parent_id: ArrayId) -> &[&'static dyn DynParentKernel] { + let Some(dense) = self.dense else { + return &[]; + }; + let dense = dense.get_or_init(|| { + let Some(max_idx) = self.keyed.iter().map(|entry| entry.parent_id.index()).max() else { + let empty: Vec>> = Vec::new(); + return empty.into_boxed_slice(); + }; + + let mut dense = (0..=max_idx).map(|_| Vec::new()).collect::>(); + for entry in self.keyed { + dense[entry.parent_id.index()].push(entry.kernel); + } + dense.into_boxed_slice() + }); + + dense + .get(parent_id.index()) + .map(Vec::as_slice) + .unwrap_or(&[]) + } /// Create a new parent kernel set with the given kernels. /// /// Use [`ParentKernelSet::lift`] to lift static rules into dynamic trait objects. pub const fn new(kernels: &'static [&'static dyn DynParentKernel]) -> Self { - Self { kernels } + Self { + keyed: &[], + dense: None, + dynamic: kernels, + } } + pub(crate) const fn new_indexed( + keyed: &'static [ParentKernelEntry], + dense: &'static ParentKernelDense, + dynamic: &'static [&'static dyn DynParentKernel], + ) -> Self { + Self { + keyed, + dense: Some(dense), + dynamic, + } + } /// Lift the given rule into a dynamic trait object. pub const fn lift>( kernel: &'static K, @@ -56,6 +106,16 @@ impl ParentKernelSet { unsafe { &*(kernel as *const K as *const ParentKernelAdapter) } } + pub(crate) const fn lift_id>( + parent_id: CachedId, + kernel: &'static K, + ) -> ParentKernelEntry { + ParentKernelEntry { + parent_id, + kernel: Self::lift(kernel), + } + } + /// Evaluate the parent kernels on the given child and parent arrays. pub fn execute( &self, @@ -64,7 +124,13 @@ impl ParentKernelSet { child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - for kernel in self.kernels.iter() { + for kernel in self.kernels_for(parent.encoding_id()) { + if let Some(reduced) = kernel.execute_parent(child, parent, child_idx, ctx)? { + return Ok(Some(reduced)); + } + } + + for kernel in self.dynamic.iter() { if !kernel.matches(parent) { continue; } diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index e505b21a199..3e5ba173fc5 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -21,9 +21,12 @@ use std::any::type_name; use std::fmt::Debug; use std::marker::PhantomData; +use std::sync::OnceLock; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use crate::ArrayId; use crate::ArrayRef; use crate::array::ArrayView; use crate::array::VTable; @@ -142,17 +145,64 @@ impl ReduceRuleSet { /// A set of parent reduction rules for a specific child array encoding. pub struct ParentRuleSet { - rules: &'static [&'static dyn DynArrayParentReduceRule], + keyed: &'static [ParentRuleEntry], + dense: Option<&'static ParentRuleDense>, + dynamic: &'static [&'static dyn DynArrayParentReduceRule], } +/// A parent reduction rule keyed by exact parent encoding id. +pub(crate) struct ParentRuleEntry { + parent_id: CachedId, + rule: &'static dyn DynArrayParentReduceRule, +} + +pub(crate) type ParentRuleDense = OnceLock>]>>; + impl ParentRuleSet { + fn rules_for(&self, parent_id: ArrayId) -> &[&'static dyn DynArrayParentReduceRule] { + let Some(dense) = self.dense else { + return &[]; + }; + let dense = dense.get_or_init(|| { + let Some(max_idx) = self.keyed.iter().map(|entry| entry.parent_id.index()).max() else { + let empty: Vec>> = Vec::new(); + return empty.into_boxed_slice(); + }; + + let mut dense = (0..=max_idx).map(|_| Vec::new()).collect::>(); + for entry in self.keyed { + dense[entry.parent_id.index()].push(entry.rule); + } + dense.into_boxed_slice() + }); + + dense + .get(parent_id.index()) + .map(Vec::as_slice) + .unwrap_or(&[]) + } /// Create a new parent rule set with the given rules. /// /// Use [`ParentRuleSet::lift`] to lift static rules into dynamic trait objects. pub const fn new(rules: &'static [&'static dyn DynArrayParentReduceRule]) -> Self { - Self { rules } + Self { + keyed: &[], + dense: None, + dynamic: rules, + } } + pub(crate) const fn new_indexed( + keyed: &'static [ParentRuleEntry], + dense: &'static ParentRuleDense, + dynamic: &'static [&'static dyn DynArrayParentReduceRule], + ) -> Self { + Self { + keyed, + dense: Some(dense), + dynamic, + } + } /// Lift the given rule into a dynamic trait object. pub const fn lift>( rule: &'static R, @@ -167,6 +217,16 @@ impl ParentRuleSet { unsafe { &*(rule as *const R as *const ParentReduceRuleAdapter) } } + pub(crate) const fn lift_id>( + parent_id: CachedId, + rule: &'static R, + ) -> ParentRuleEntry { + ParentRuleEntry { + parent_id, + rule: Self::lift(rule), + } + } + /// Evaluate the parent reduction rules on the given child and parent arrays. pub fn evaluate( &self, @@ -174,7 +234,32 @@ impl ParentRuleSet { parent: &ArrayRef, child_idx: usize, ) -> VortexResult> { - for rule in self.rules.iter() { + for rule in self.rules_for(parent.encoding_id()) { + if let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? { + // Debug assertions because these checks are already run elsewhere. + #[cfg(debug_assertions)] + { + vortex_error::vortex_ensure!( + reduced.len() == parent.len(), + "Reduced array length mismatch from {:?}\nFrom:\n{}\nTo:\n{}", + rule, + parent.encoding_id(), + reduced.encoding_id() + ); + vortex_error::vortex_ensure!( + reduced.dtype() == parent.dtype(), + "Reduced array dtype mismatch from {:?}\nFrom:\n{}\nTo:\n{}", + rule, + parent.encoding_id(), + reduced.encoding_id() + ); + } + + return Ok(Some(reduced)); + } + } + + for rule in self.dynamic.iter() { if !rule.matches(parent) { continue; } diff --git a/vortex-session/public-api.lock b/vortex-session/public-api.lock index a63e955f77a..16af357ecf5 100644 --- a/vortex-session/public-api.lock +++ b/vortex-session/public-api.lock @@ -46,6 +46,8 @@ impl vortex_session::registry::Id pub fn vortex_session::registry::Id::as_str(&self) -> &str +pub fn vortex_session::registry::Id::index(&self) -> usize + pub fn vortex_session::registry::Id::new(s: &str) -> Self pub fn vortex_session::registry::Id::new_static(s: &'static str) -> Self diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs index a739b9fdda3..2a19754d06a 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -14,6 +14,7 @@ use std::sync::Arc; use std::sync::LazyLock; use std::sync::OnceLock; +use lasso::Key; use lasso::Spur; use lasso::ThreadedRodeo; use parking_lot::RwLock; @@ -49,6 +50,11 @@ impl Id { // pointers are stable for the lifetime of the program. unsafe { &*(s as *const str) } } + + /// Returns this ID's dense index in the global interner. + pub fn index(&self) -> usize { + self.0.into_usize() + } } impl From<&str> for Id { From 3f38b25996741497794fde4698da54045850823d Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 13:03:58 -0400 Subject: [PATCH 04/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index cd6e63e065f..55a69157ddf 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -23,7 +23,6 @@ use std::fmt::Display; use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; From a518164aa79e9966dbaf93c6e2b42a16a05eef39 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 13:15:05 -0400 Subject: [PATCH 05/16] fix Signed-off-by: Joe Isaacs --- vortex-array/benches/chunked_dict_builder.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index 85c8737356f..77583bbaa56 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -60,6 +60,11 @@ fn chunked_dict_primitive_into_canonical( { let chunk = gen_dict_primitive_chunks::(len, unique_values, chunk_count); + chunk + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + bencher .with_inputs(|| (chunk.clone(), SESSION.create_execution_ctx())) .bench_values(|(chunk, mut ctx)| chunk.execute::(&mut ctx)) From b6e94e0e56657ef22c51f2980fe39e9147f822a9 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 13:19:41 -0400 Subject: [PATCH 06/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 55a69157ddf..631a32c8db3 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -46,7 +46,7 @@ pub(crate) fn max_iterations() -> usize { Ok(val) => val.parse::().unwrap_or_else(|e| { vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid usize: {e}") }), - Err(VarError::NotPresent) => 128, + Err(VarError::NotPresent) => 1024, Err(VarError::NotUnicode(_)) => { vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid unicode string") } From c2800f825242e70474295762dc677c1902aeb233 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 13:35:02 -0400 Subject: [PATCH 07/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 631a32c8db3..3f48c01f7c7 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -46,7 +46,7 @@ pub(crate) fn max_iterations() -> usize { Ok(val) => val.parse::().unwrap_or_else(|e| { vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid usize: {e}") }), - Err(VarError::NotPresent) => 1024, + Err(VarError::NotPresent) => 1025, Err(VarError::NotUnicode(_)) => { vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid unicode string") } From 7e5c9feafd938625dc593f421bdf6adb6b869095 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 16:02:19 -0400 Subject: [PATCH 08/16] fix Signed-off-by: Joe Isaacs --- encodings/alp/src/alp/rules.rs | 41 +++++++++++++------ encodings/alp/src/alp_rd/kernel.rs | 18 +++++--- encodings/alp/src/alp_rd/rules.rs | 16 ++++++-- encodings/bytebool/src/kernel.rs | 14 ++++++- encodings/bytebool/src/rules.rs | 18 +++++--- .../datetime-parts/src/compute/kernel.rs | 22 ++++++++-- encodings/datetime-parts/src/compute/rules.rs | 41 +++++++++++++++---- .../src/decimal_byte_parts/compute/kernel.rs | 22 ++++++++-- .../src/decimal_byte_parts/rules.rs | 36 ++++++++++++---- .../src/bitpacking/vtable/kernels.rs | 19 +++++++-- .../fastlanes/src/bitpacking/vtable/rules.rs | 19 +++++++-- encodings/fastlanes/src/delta/vtable/rules.rs | 16 ++++++-- encodings/fastlanes/src/for/vtable/kernels.rs | 16 ++++++-- encodings/fastlanes/src/for/vtable/rules.rs | 26 ++++++++---- encodings/fastlanes/src/rle/kernel.rs | 12 +++++- encodings/fastlanes/src/rle/vtable/rules.rs | 14 ++++++- encodings/fsst/src/kernel.rs | 20 ++++++--- encodings/fsst/src/rules.rs | 16 ++++++-- encodings/parquet-variant/src/kernel.rs | 27 +++++++++--- encodings/pco/src/rules.rs | 16 ++++++-- encodings/runend/src/kernel.rs | 30 ++++++++++---- encodings/runend/src/rules.rs | 23 ++++++++--- encodings/sequence/src/kernel.rs | 24 ++++++++--- encodings/sequence/src/rules.rs | 21 +++++++--- encodings/sparse/src/kernel.rs | 21 +++++++--- encodings/sparse/src/rules.rs | 16 ++++++-- encodings/zigzag/src/kernel.rs | 14 ++++++- encodings/zigzag/src/rules.rs | 20 ++++++--- encodings/zstd/src/rules.rs | 16 ++++++-- vortex-array/benches/chunked_dict_builder.rs | 2 +- vortex-array/src/arrays/bool/compute/rules.rs | 22 ++++++---- vortex-array/src/arrays/bool/vtable/kernel.rs | 20 +++++++-- .../src/arrays/constant/compute/rules.rs | 38 ++++++++++++----- .../src/arrays/constant/compute/take.rs | 19 +++++++-- .../src/arrays/decimal/compute/rules.rs | 18 +++++--- .../src/arrays/decimal/vtable/kernel.rs | 27 +++++++++--- .../src/arrays/extension/compute/rules.rs | 28 +++++++++---- .../src/arrays/extension/vtable/kernel.rs | 20 +++++++-- vortex-array/src/arrays/filter/rules.rs | 14 ++++++- .../arrays/fixed_size_list/compute/rules.rs | 28 ++++++++++--- .../arrays/fixed_size_list/vtable/kernel.rs | 17 ++++++-- .../src/arrays/fixed_size_list/vtable/mod.rs | 2 +- .../src/arrays/list/compute/kernels.rs | 17 ++++++-- vortex-array/src/arrays/list/compute/rules.rs | 19 ++++++--- .../src/arrays/listview/compute/rules.rs | 22 ++++++---- .../src/arrays/masked/compute/rules.rs | 21 +++++++--- vortex-array/src/arrays/null/compute/rules.rs | 23 +++++++---- vortex-array/src/arrays/null/compute/take.rs | 19 +++++++-- .../src/arrays/patched/compute/rules.rs | 20 +++++++-- .../src/arrays/patched/vtable/kernels.rs | 20 +++++++-- vortex-array/src/arrays/scalar_fn/rules.rs | 19 +++++++-- vortex-array/src/arrays/slice/rules.rs | 15 ++++++- .../src/arrays/struct_/compute/rules.rs | 22 ++++++---- .../src/arrays/struct_/vtable/kernel.rs | 17 ++++++-- .../src/arrays/varbin/compute/rules.rs | 19 ++++++--- .../src/arrays/varbin/vtable/kernel.rs | 25 ++++++++--- .../src/arrays/varbinview/compute/rules.rs | 23 ++++++++--- .../src/arrays/varbinview/vtable/kernel.rs | 20 +++++++-- vortex-array/src/kernel.rs | 8 ++-- vortex-array/src/optimizer/rules.rs | 8 ++-- 60 files changed, 921 insertions(+), 285 deletions(-) diff --git a/encodings/alp/src/alp/rules.rs b/encodings/alp/src/alp/rules.rs index fe13682ccd8..a3d068f132c 100644 --- a/encodings/alp/src/alp/rules.rs +++ b/encodings/alp/src/alp/rules.rs @@ -4,26 +4,41 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::registry::CachedId; use crate::ALP; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(ALP)), - ParentKernelSet::lift(&FilterExecuteAdaptor(ALP)), - ParentKernelSet::lift(&MaskExecuteAdaptor(ALP)), - ParentKernelSet::lift(&SliceExecuteAdaptor(ALP)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ALP)), -]); - -pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)), - ParentRuleSet::lift(&CastReduceAdaptor(ALP)), - ParentRuleSet::lift(&MaskReduceAdaptor(ALP)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 5] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.binary"), &CompareExecuteAdaptor(ALP)), + ParentKernelSet::lift_id(CachedId::new("vortex.filter"), &FilterExecuteAdaptor(ALP)), + ParentKernelSet::lift_id(CachedId::new("vortex.mask"), &MaskExecuteAdaptor(ALP)), + ParentKernelSet::lift_id(CachedId::new("vortex.slice"), &SliceExecuteAdaptor(ALP)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(ALP)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); + +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.between"), &BetweenReduceAdaptor(ALP)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(ALP)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(ALP)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/alp/src/alp_rd/kernel.rs b/encodings/alp/src/alp_rd/kernel.rs index 66721cbb1f9..b560d2b7e9a 100644 --- a/encodings/alp/src/alp_rd/kernel.rs +++ b/encodings/alp/src/alp_rd/kernel.rs @@ -4,12 +4,20 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_session::registry::CachedId; use crate::alp_rd::ALPRD; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&SliceExecuteAdaptor(ALPRD)), - ParentKernelSet::lift(&FilterExecuteAdaptor(ALPRD)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ALPRD)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.slice"), &SliceExecuteAdaptor(ALPRD)), + ParentKernelSet::lift_id(CachedId::new("vortex.filter"), &FilterExecuteAdaptor(ALPRD)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(ALPRD)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/alp/src/alp_rd/rules.rs b/encodings/alp/src/alp_rd/rules.rs index be5eda6e7f3..884b2b350a7 100644 --- a/encodings/alp/src/alp_rd/rules.rs +++ b/encodings/alp/src/alp_rd/rules.rs @@ -1,13 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::registry::CachedId; use crate::alp_rd::ALPRD; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(ALPRD)), - ParentRuleSet::lift(&MaskReduceAdaptor(ALPRD)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(ALPRD)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(ALPRD)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index 4373520039a..872ef6fa7cc 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -2,9 +2,19 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_session::registry::CachedId; use crate::ByteBool; -pub(crate) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool))]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 1] = [ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(ByteBool), +)]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/bytebool/src/rules.rs b/encodings/bytebool/src/rules.rs index f67d3567326..710ae3b5e44 100644 --- a/encodings/bytebool/src/rules.rs +++ b/encodings/bytebool/src/rules.rs @@ -2,14 +2,22 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::registry::CachedId; use crate::ByteBool; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(ByteBool)), - ParentRuleSet::lift(&MaskReduceAdaptor(ByteBool)), - ParentRuleSet::lift(&SliceReduceAdaptor(ByteBool)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(ByteBool)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(ByteBool)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(ByteBool)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/datetime-parts/src/compute/kernel.rs b/encodings/datetime-parts/src/compute/kernel.rs index e07f1d00f38..d1868f7dd0b 100644 --- a/encodings/datetime-parts/src/compute/kernel.rs +++ b/encodings/datetime-parts/src/compute/kernel.rs @@ -2,12 +2,26 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::registry::CachedId; use crate::DateTimeParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DateTimeParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DateTimeParts)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(DateTimeParts), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(DateTimeParts), + ), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index c18d959cd7e..85dd5c78024 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -18,6 +18,8 @@ use vortex_array::dtype::DType; use vortex_array::extension::datetime::Timestamp; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::optimizer::rules::ArrayParentReduceRule; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::binary::Binary; @@ -25,18 +27,41 @@ use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::DateTimeParts; use crate::array::DateTimePartsArrayExt; use crate::timestamp; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&DTPFilterPushDownRule), - ParentRuleSet::lift(&DTPComparisonPushDownRule), - ParentRuleSet::lift(&CastReduceAdaptor(DateTimeParts)), - ParentRuleSet::lift(&FilterReduceAdaptor(DateTimeParts)), - ParentRuleSet::lift(&MaskReduceAdaptor(DateTimeParts)), - ParentRuleSet::lift(&SliceReduceAdaptor(DateTimeParts)), -]); + +static KEYED_PARENT_RULES: [ParentRuleEntry; 4] = [ + ParentRuleSet::lift_id( + CachedId::new("vortex.cast"), + &CastReduceAdaptor(DateTimeParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterReduceAdaptor(DateTimeParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.mask"), + &MaskReduceAdaptor(DateTimeParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(DateTimeParts), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ + ParentRuleSet::lift(&DTPFilterPushDownRule), + ParentRuleSet::lift(&DTPComparisonPushDownRule), + ], +); /// Push the filter into the days column of a date time parts, we could extend this to other fields /// but its less clear if that is beneficial. diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 9e16b4f81dc..6b36a7e2146 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -2,12 +2,26 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::registry::CachedId; use crate::DecimalByteParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DecimalByteParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DecimalByteParts)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(DecimalByteParts), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(DecimalByteParts), + ), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs index 7fafa0d5289..458fe40c585 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs @@ -8,22 +8,44 @@ use vortex_array::arrays::Filter; use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ArrayParentReduceRule; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArrayExt; -pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule), - ParentRuleSet::lift(&CastReduceAdaptor(DecimalByteParts)), - ParentRuleSet::lift(&FilterReduceAdaptor(DecimalByteParts)), - ParentRuleSet::lift(&MaskReduceAdaptor(DecimalByteParts)), - ParentRuleSet::lift(&SliceReduceAdaptor(DecimalByteParts)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 4] = [ + ParentRuleSet::lift_id( + CachedId::new("vortex.cast"), + &CastReduceAdaptor(DecimalByteParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterReduceAdaptor(DecimalByteParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.mask"), + &MaskReduceAdaptor(DecimalByteParts), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(DecimalByteParts), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule)], +); #[derive(Debug)] struct DecimalBytePartsFilterPushDownRule; diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index 128ff77d99c..df95a29e3ff 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -3,11 +3,22 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_session::registry::CachedId; use crate::BitPacked; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&TakeExecuteAdaptor(BitPacked)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(BitPacked), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(BitPacked)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/fastlanes/src/bitpacking/vtable/rules.rs b/encodings/fastlanes/src/bitpacking/vtable/rules.rs index a26989569b2..ac1aa3c57a0 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/rules.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/rules.rs @@ -2,12 +2,23 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::BitPacked; -pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(BitPacked)), - ParentRuleSet::lift(&SliceReduceAdaptor(BitPacked)), -]); +static KEYED_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(BitPacked)), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(BitPacked), + ), +]; + +static KEYED_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_RULES, &KEYED_RULES_DENSE, &[]); diff --git a/encodings/fastlanes/src/delta/vtable/rules.rs b/encodings/fastlanes/src/delta/vtable/rules.rs index d6892897ab5..7ede228507d 100644 --- a/encodings/fastlanes/src/delta/vtable/rules.rs +++ b/encodings/fastlanes/src/delta/vtable/rules.rs @@ -2,12 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::delta::vtable::Delta; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&SliceReduceAdaptor(Delta)), - ParentRuleSet::lift(&CastReduceAdaptor(Delta)), -]); +static KEYED_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Delta)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Delta)), +]; + +static KEYED_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_RULES, &KEYED_RULES_DENSE, &[]); diff --git a/encodings/fastlanes/src/for/vtable/kernels.rs b/encodings/fastlanes/src/for/vtable/kernels.rs index 331b13eceef..39caf6440da 100644 --- a/encodings/fastlanes/src/for/vtable/kernels.rs +++ b/encodings/fastlanes/src/for/vtable/kernels.rs @@ -2,12 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::registry::CachedId; use crate::FoR; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(FoR)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FoR)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.binary"), &CompareExecuteAdaptor(FoR)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(FoR)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/fastlanes/src/for/vtable/rules.rs b/encodings/fastlanes/src/for/vtable/rules.rs index 8d220a86d7f..236033ae03b 100644 --- a/encodings/fastlanes/src/for/vtable/rules.rs +++ b/encodings/fastlanes/src/for/vtable/rules.rs @@ -8,20 +8,32 @@ use vortex_array::arrays::Filter; use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ArrayParentReduceRule; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::FoR; use crate::r#for::array::FoRArrayExt; -pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - // TODO: add BetweenReduceAdaptor(FoR) - ParentRuleSet::lift(&FoRFilterPushDownRule), - ParentRuleSet::lift(&FilterReduceAdaptor(FoR)), - ParentRuleSet::lift(&SliceReduceAdaptor(FoR)), - ParentRuleSet::lift(&CastReduceAdaptor(FoR)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(FoR)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(FoR)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(FoR)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ + // TODO: add BetweenReduceAdaptor(FoR) + ParentRuleSet::lift(&FoRFilterPushDownRule), + ], +); #[derive(Debug)] struct FoRFilterPushDownRule; diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index 14eaf6dda3d..738175cb4e3 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -9,15 +9,25 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 1] = [ParentKernelSet::lift_id( + CachedId::new("vortex.slice"), + &SliceExecuteAdaptor(RLE), +)]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + pub(crate) static PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&SliceExecuteAdaptor(RLE))]); + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); impl SliceKernel for RLE { fn slice( diff --git a/encodings/fastlanes/src/rle/vtable/rules.rs b/encodings/fastlanes/src/rle/vtable/rules.rs index 50dd1f16a5a..c11678ff044 100644 --- a/encodings/fastlanes/src/rle/vtable/rules.rs +++ b/encodings/fastlanes/src/rle/vtable/rules.rs @@ -1,10 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::RLE; -pub(crate) const RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&CastReduceAdaptor(RLE))]); +static KEYED_RULES: [ParentRuleEntry; 1] = [ParentRuleSet::lift_id( + CachedId::new("vortex.cast"), + &CastReduceAdaptor(RLE), +)]; + +static KEYED_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_RULES, &KEYED_RULES_DENSE, &[]); diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 05b58e3215d..19568bddc00 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -3,18 +3,26 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_session::registry::CachedId; use crate::FSST; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(FSST)), - ParentKernelSet::lift(&FilterExecuteAdaptor(FSST)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FSST)), - ParentKernelSet::lift(&LikeExecuteAdaptor(FSST)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 4] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.binary"), &CompareExecuteAdaptor(FSST)), + ParentKernelSet::lift_id(CachedId::new("vortex.filter"), &FilterExecuteAdaptor(FSST)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(FSST)), + ParentKernelSet::lift_id(CachedId::new("vortex.like"), &LikeExecuteAdaptor(FSST)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); #[cfg(test)] mod tests { diff --git a/encodings/fsst/src/rules.rs b/encodings/fsst/src/rules.rs index bec81c0fd3a..d2494edfb6f 100644 --- a/encodings/fsst/src/rules.rs +++ b/encodings/fsst/src/rules.rs @@ -2,12 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::FSST; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&SliceReduceAdaptor(FSST)), - ParentRuleSet::lift(&CastReduceAdaptor(FSST)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(FSST)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(FSST)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 4743b9fc1a7..c17068fe915 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -13,18 +13,35 @@ use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::filter::FilterKernel; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::registry::CachedId; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(ParquetVariant)), - ParentKernelSet::lift(&SliceExecuteAdaptor(ParquetVariant)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ParquetVariant)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(ParquetVariant), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.slice"), + &SliceExecuteAdaptor(ParquetVariant), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(ParquetVariant), + ), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); impl SliceKernel for ParquetVariant { fn slice( diff --git a/encodings/pco/src/rules.rs b/encodings/pco/src/rules.rs index 4d27392dc01..7882a9d9b66 100644 --- a/encodings/pco/src/rules.rs +++ b/encodings/pco/src/rules.rs @@ -2,12 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::Pco; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&SliceReduceAdaptor(Pco)), - ParentRuleSet::lift(&CastReduceAdaptor(Pco)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Pco)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Pco)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 9480e9e6f04..dea1772d0a5 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -14,21 +14,37 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&RunEndSliceKernel), - ParentKernelSet::lift(&FilterExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&TakeExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&RunEndTakeFrom), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 4] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(RunEnd), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.slice"), &RunEndSliceKernel), + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(RunEnd), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(RunEnd)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new_indexed( + &KEYED_PARENT_KERNELS, + &KEYED_PARENT_KERNELS_DENSE, + &[ParentKernelSet::lift(&RunEndTakeFrom)], +); /// Kernel to execute slicing on a RunEnd array. /// diff --git a/encodings/runend/src/rules.rs b/encodings/runend/src/rules.rs index 73f9ef3cb7a..3c05d4307ca 100644 --- a/encodings/runend/src/rules.rs +++ b/encodings/runend/src/rules.rs @@ -12,22 +12,35 @@ use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnVTable; use vortex_array::dtype::DType; use vortex_array::optimizer::rules::ArrayParentReduceRule; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::fill_null::FillNullReduceAdaptor; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::RunEnd; use crate::array::RunEndArrayExt; -pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ // CastReduceAdaptor must come before RunEndScalarFnRule so that cast operations are executed // eagerly (surfacing out-of-range errors immediately) rather than being pushed lazily into // the values array by the generic scalar function push-down rule. - ParentRuleSet::lift(&CastReduceAdaptor(RunEnd)), - ParentRuleSet::lift(&RunEndScalarFnRule), - ParentRuleSet::lift(&FillNullReduceAdaptor(RunEnd)), -]); + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(RunEnd)), + ParentRuleSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullReduceAdaptor(RunEnd), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ParentRuleSet::lift(&RunEndScalarFnRule)], +); /// A rule to push down scalar functions through run-end encoding into the values array. /// diff --git a/encodings/sequence/src/kernel.rs b/encodings/sequence/src/kernel.rs index 394a52ec69e..c145e6b537d 100644 --- a/encodings/sequence/src/kernel.rs +++ b/encodings/sequence/src/kernel.rs @@ -3,13 +3,27 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::registry::CachedId; use crate::Sequence; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&FilterExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sequence)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(Sequence), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(Sequence), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Sequence)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/sequence/src/rules.rs b/encodings/sequence/src/rules.rs index f87c08e8109..af2cc07b3e8 100644 --- a/encodings/sequence/src/rules.rs +++ b/encodings/sequence/src/rules.rs @@ -2,14 +2,25 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor; +use vortex_session::registry::CachedId; use crate::Sequence; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(Sequence)), - ParentRuleSet::lift(&ListContainsElementReduceAdaptor(Sequence)), - ParentRuleSet::lift(&SliceReduceAdaptor(Sequence)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Sequence)), + ParentRuleSet::lift_id( + CachedId::new("vortex.list_contains"), + &ListContainsElementReduceAdaptor(Sequence), + ), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Sequence)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/sparse/src/kernel.rs b/encodings/sparse/src/kernel.rs index 18928ea0142..e570612fcaa 100644 --- a/encodings/sparse/src/kernel.rs +++ b/encodings/sparse/src/kernel.rs @@ -4,12 +4,23 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_session::registry::CachedId; use crate::Sparse; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&SliceExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sparse)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(Sparse), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.slice"), &SliceExecuteAdaptor(Sparse)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Sparse)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/sparse/src/rules.rs b/encodings/sparse/src/rules.rs index c53e7e4c6a5..378c4bf556a 100644 --- a/encodings/sparse/src/rules.rs +++ b/encodings/sparse/src/rules.rs @@ -5,18 +5,26 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::not::NotReduce; use vortex_array::scalar_fn::fns::not::NotReduceAdaptor; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::Sparse; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(Sparse)), - ParentRuleSet::lift(&NotReduceAdaptor(Sparse)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Sparse)), + ParentRuleSet::lift_id(CachedId::new("vortex.not"), &NotReduceAdaptor(Sparse)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); impl NotReduce for Sparse { fn invert(array: ArrayView<'_, Self>) -> VortexResult> { diff --git a/encodings/zigzag/src/kernel.rs b/encodings/zigzag/src/kernel.rs index d0096abaae1..69aaed4ece6 100644 --- a/encodings/zigzag/src/kernel.rs +++ b/encodings/zigzag/src/kernel.rs @@ -2,9 +2,19 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::kernel::ParentKernelDense; +use vortex_array::kernel::ParentKernelEntry; use vortex_array::kernel::ParentKernelSet; +use vortex_session::registry::CachedId; use crate::ZigZag; -pub(crate) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(ZigZag))]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 1] = [ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(ZigZag), +)]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/encodings/zigzag/src/rules.rs b/encodings/zigzag/src/rules.rs index c3b612d101d..601deb25c9d 100644 --- a/encodings/zigzag/src/rules.rs +++ b/encodings/zigzag/src/rules.rs @@ -3,15 +3,23 @@ use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::registry::CachedId; use crate::ZigZag; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(ZigZag)), - ParentRuleSet::lift(&FilterReduceAdaptor(ZigZag)), - ParentRuleSet::lift(&MaskReduceAdaptor(ZigZag)), - ParentRuleSet::lift(&SliceReduceAdaptor(ZigZag)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 4] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(ZigZag)), + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(ZigZag)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(ZigZag)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(ZigZag)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/encodings/zstd/src/rules.rs b/encodings/zstd/src/rules.rs index cfbf93b09e8..af93689acc7 100644 --- a/encodings/zstd/src/rules.rs +++ b/encodings/zstd/src/rules.rs @@ -2,12 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleDense; +use vortex_array::optimizer::rules::ParentRuleEntry; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_session::registry::CachedId; use crate::Zstd; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&SliceReduceAdaptor(Zstd)), - ParentRuleSet::lift(&CastReduceAdaptor(Zstd)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Zstd)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Zstd)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index 77583bbaa56..3256d2201a4 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -63,7 +63,7 @@ fn chunked_dict_primitive_into_canonical( chunk .clone() .execute::(&mut SESSION.create_execution_ctx()) - .unwrap(); + .vortex_expect("warmup execute failed"); bencher .with_inputs(|| (chunk.clone(), SESSION.create_execution_ctx())) diff --git a/vortex-array/src/arrays/bool/compute/rules.rs b/vortex-array/src/arrays/bool/compute/rules.rs index ed2e7a4dc28..23a15fc9f4c 100644 --- a/vortex-array/src/arrays/bool/compute/rules.rs +++ b/vortex-array/src/arrays/bool/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -13,17 +14,24 @@ use crate::arrays::bool::BoolArrayExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&BoolMaskedValidityRule), - ParentRuleSet::lift(&CastReduceAdaptor(Bool)), - ParentRuleSet::lift(&MaskReduceAdaptor(Bool)), - ParentRuleSet::lift(&SliceReduceAdaptor(Bool)), - ParentRuleSet::lift(&FilterReduceAdaptor(Bool)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.masked"), &BoolMaskedValidityRule), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Bool)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Bool)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Bool)), + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(Bool)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Rule to push down validity masking from MaskedArray parent into BoolArray child. /// diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index e8fa5f30e9c..f87b71f88d3 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Bool; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FillNullExecuteAdaptor(Bool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Bool)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullExecuteAdaptor(Bool), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Bool)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/constant/compute/rules.rs b/vortex-array/src/arrays/constant/compute/rules.rs index f97e6c596c7..c10346b1917 100644 --- a/vortex-array/src/arrays/constant/compute/rules.rs +++ b/vortex-array/src/arrays/constant/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -13,22 +14,39 @@ use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::between::BetweenReduceAdaptor; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor; use crate::scalar_fn::fns::not::NotReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&BetweenReduceAdaptor(Constant)), - ParentRuleSet::lift(&CastReduceAdaptor(Constant)), - ParentRuleSet::lift(&ConstantFilterRule), - ParentRuleSet::lift(&FillNullReduceAdaptor(Constant)), - ParentRuleSet::lift(&FilterReduceAdaptor(Constant)), - ParentRuleSet::lift(&NotReduceAdaptor(Constant)), - ParentRuleSet::lift(&SliceReduceAdaptor(Constant)), - ParentRuleSet::lift(&TakeReduceAdaptor(Constant)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 6] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Constant)), + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &ConstantFilterRule), + ParentRuleSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullReduceAdaptor(Constant), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterReduceAdaptor(Constant), + ), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Constant)), + ParentRuleSet::lift_id(CachedId::new("vortex.dict"), &TakeReduceAdaptor(Constant)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = ParentRuleSet::new_indexed( + &KEYED_PARENT_RULES, + &KEYED_PARENT_RULES_DENSE, + &[ + ParentRuleSet::lift(&BetweenReduceAdaptor(Constant)), + ParentRuleSet::lift(&NotReduceAdaptor(Constant)), + ], +); #[derive(Debug)] struct ConstantFilterRule; diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 15e67a02929..6951288bb03 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -3,6 +3,7 @@ use vortex_error::VortexResult; use vortex_mask::AllOr; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -14,6 +15,8 @@ use crate::arrays::ConstantArray; use crate::arrays::MaskedArray; use crate::arrays::dict::TakeReduce; use crate::arrays::dict::TakeReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar::Scalar; use crate::validity::Validity; @@ -59,10 +62,18 @@ impl TakeReduce for Constant { } } -impl Constant { - pub const TAKE_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&TakeReduceAdaptor::(Self))]); -} +#[allow(dead_code)] +static KEYED_TAKE_RULES: [ParentRuleEntry; 1] = [ParentRuleSet::lift_id( + CachedId::new("vortex.dict"), + &TakeReduceAdaptor::(Constant), +)]; + +#[allow(dead_code)] +static KEYED_TAKE_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +#[allow(dead_code)] +pub(crate) static TAKE_PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_TAKE_RULES, &KEYED_TAKE_RULES_DENSE, &[]); #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/decimal/compute/rules.rs b/vortex-array/src/arrays/decimal/compute/rules.rs index fae0c3dd866..5d1f770713b 100644 --- a/vortex-array/src/arrays/decimal/compute/rules.rs +++ b/vortex-array/src/arrays/decimal/compute/rules.rs @@ -4,6 +4,7 @@ use std::ops::Range; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -15,14 +16,21 @@ use crate::arrays::slice::SliceReduce; use crate::arrays::slice::SliceReduceAdaptor; use crate::match_each_decimal_value_type; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&DecimalMaskedValidityRule), - ParentRuleSet::lift(&MaskReduceAdaptor(Decimal)), - ParentRuleSet::lift(&SliceReduceAdaptor(Decimal)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.masked"), &DecimalMaskedValidityRule), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Decimal)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Decimal)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Rule to push down validity masking from MaskedArray parent into DecimalArray child. /// diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 9ad3803be1d..4161725cbed 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -1,16 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Decimal; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&CastExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Decimal)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 4] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.between"), + &BetweenExecuteAdaptor(Decimal), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.cast"), &CastExecuteAdaptor(Decimal)), + ParentKernelSet::lift_id( + CachedId::new("vortex.fill_null"), + &FillNullExecuteAdaptor(Decimal), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Decimal)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 568fe1c061c..a2515d8aa66 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -16,6 +17,8 @@ use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; use crate::scalar::Scalar; @@ -44,13 +47,24 @@ impl ArrayReduceRule for ExtensionConstantRule { } } -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&ExtensionFilterPushDownRule), - ParentRuleSet::lift(&CastReduceAdaptor(Extension)), - ParentRuleSet::lift(&FilterReduceAdaptor(Extension)), - ParentRuleSet::lift(&MaskReduceAdaptor(Extension)), - ParentRuleSet::lift(&SliceReduceAdaptor(Extension)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &ExtensionFilterPushDownRule), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Extension)), + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterReduceAdaptor(Extension), + ), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Extension)), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(Extension), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Push filter operations into the storage array of an extension array. #[derive(Debug)] diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index 4cc5f9eb4ac..65dd99b231c 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Extension; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Extension)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Extension)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(Extension), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Extension)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/filter/rules.rs b/vortex-array/src/arrays/filter/rules.rs index 45e938d7768..0996097ab6d 100644 --- a/vortex-array/src/arrays/filter/rules.rs +++ b/vortex-array/src/arrays/filter/rules.rs @@ -3,6 +3,7 @@ use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Canonical; @@ -15,11 +16,20 @@ use crate::arrays::filter::FilterArrayExt; use crate::arrays::struct_::StructDataParts; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; -pub(super) const PARENT_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&FilterFilterRule)]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 1] = [ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterFilterRule, +)]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&TrivialFilterRule, &FilterStructRule]); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/rules.rs b/vortex-array/src/arrays/fixed_size_list/compute/rules.rs index da1d91423e2..4686b907c0f 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/rules.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/rules.rs @@ -1,14 +1,32 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::FixedSizeList; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(FixedSizeList)), - ParentRuleSet::lift(&MaskReduceAdaptor(FixedSizeList)), - ParentRuleSet::lift(&SliceReduceAdaptor(FixedSizeList)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id( + CachedId::new("vortex.cast"), + &CastReduceAdaptor(FixedSizeList), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.mask"), + &MaskReduceAdaptor(FixedSizeList), + ), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(FixedSizeList), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs index c28a3a2a805..eb4e5d5e6d6 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -1,11 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::FixedSizeList; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; -impl FixedSizeList { - pub(crate) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(FixedSizeList))]); -} +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 1] = [ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(FixedSizeList), +)]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index a3724b06fca..d8f896a582e 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -91,7 +91,7 @@ impl VTable for FixedSizeList { child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - Self::PARENT_KERNELS.execute(array, parent, child_idx, ctx) + kernel::PARENT_KERNELS.execute(array, parent, child_idx, ctx) } fn serialize( diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index b4268642e82..75335153ae1 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -1,12 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::List; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(List)), - ParentKernelSet::lift(&TakeExecuteAdaptor(List)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.filter"), &FilterExecuteAdaptor(List)), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(List)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(crate) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/list/compute/rules.rs b/vortex-array/src/arrays/list/compute/rules.rs index d2aca066f7f..8fabcd67b10 100644 --- a/vortex-array/src/arrays/list/compute/rules.rs +++ b/vortex-array/src/arrays/list/compute/rules.rs @@ -1,14 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::List; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(List)), - ParentRuleSet::lift(&MaskReduceAdaptor(List)), - ParentRuleSet::lift(&SliceReduceAdaptor(List)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(List)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(List)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(List)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/listview/compute/rules.rs b/vortex-array/src/arrays/listview/compute/rules.rs index 9b9396d7fe7..b3ed2f9d8cf 100644 --- a/vortex-array/src/arrays/listview/compute/rules.rs +++ b/vortex-array/src/arrays/listview/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -13,17 +14,24 @@ use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&ListViewFilterPushDown), - ParentRuleSet::lift(&CastReduceAdaptor(ListView)), - ParentRuleSet::lift(&MaskReduceAdaptor(ListView)), - ParentRuleSet::lift(&SliceReduceAdaptor(ListView)), - ParentRuleSet::lift(&TakeReduceAdaptor(ListView)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &ListViewFilterPushDown), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(ListView)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(ListView)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(ListView)), + ParentRuleSet::lift_id(CachedId::new("vortex.dict"), &TakeReduceAdaptor(ListView)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); #[derive(Debug)] struct ListViewFilterPushDown; diff --git a/vortex-array/src/arrays/masked/compute/rules.rs b/vortex-array/src/arrays/masked/compute/rules.rs index 3accb455c3f..9e3c1af4d6d 100644 --- a/vortex-array/src/arrays/masked/compute/rules.rs +++ b/vortex-array/src/arrays/masked/compute/rules.rs @@ -1,16 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Masked; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&FilterReduceAdaptor(Masked)), - ParentRuleSet::lift(&MaskReduceAdaptor(Masked)), - ParentRuleSet::lift(&SliceReduceAdaptor(Masked)), - ParentRuleSet::lift(&TakeReduceAdaptor(Masked)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 4] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(Masked)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Masked)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Masked)), + ParentRuleSet::lift_id(CachedId::new("vortex.dict"), &TakeReduceAdaptor(Masked)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/null/compute/rules.rs b/vortex-array/src/arrays/null/compute/rules.rs index f63c904667a..338e1c24b2f 100644 --- a/vortex-array/src/arrays/null/compute/rules.rs +++ b/vortex-array/src/arrays/null/compute/rules.rs @@ -1,18 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Null; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&FilterReduceAdaptor(Null)), - ParentRuleSet::lift(&CastReduceAdaptor(Null)), - ParentRuleSet::lift(&MaskReduceAdaptor(Null)), - ParentRuleSet::lift(&SliceReduceAdaptor(Null)), - ParentRuleSet::lift(&TakeReduceAdaptor(Null)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.filter"), &FilterReduceAdaptor(Null)), + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(Null)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Null)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Null)), + ParentRuleSet::lift_id(CachedId::new("vortex.dict"), &TakeReduceAdaptor(Null)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/null/compute/take.rs b/vortex-array/src/arrays/null/compute/take.rs index 5352ae67653..44721b91570 100644 --- a/vortex-array/src/arrays/null/compute/take.rs +++ b/vortex-array/src/arrays/null/compute/take.rs @@ -3,6 +3,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -14,6 +15,8 @@ use crate::arrays::NullArray; use crate::arrays::dict::TakeReduce; use crate::arrays::dict::TakeReduceAdaptor; use crate::match_each_integer_ptype; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; impl TakeReduce for Null { @@ -35,7 +38,15 @@ impl TakeReduce for Null { } } -impl Null { - pub const TAKE_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&TakeReduceAdaptor::(Self))]); -} +#[allow(dead_code)] +static KEYED_TAKE_RULES: [ParentRuleEntry; 1] = [ParentRuleSet::lift_id( + CachedId::new("vortex.dict"), + &TakeReduceAdaptor::(Null), +)]; + +#[allow(dead_code)] +static KEYED_TAKE_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +#[allow(dead_code)] +pub(crate) static TAKE_PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_TAKE_RULES, &KEYED_TAKE_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/patched/compute/rules.rs b/vortex-array/src/arrays/patched/compute/rules.rs index 3ecb25c1efa..a6ed0fc19b1 100644 --- a/vortex-array/src/arrays/patched/compute/rules.rs +++ b/vortex-array/src/arrays/patched/compute/rules.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Patched; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&FilterReduceAdaptor(Patched)), - ParentRuleSet::lift(&SliceReduceAdaptor(Patched)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &FilterReduceAdaptor(Patched), + ), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Patched)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/patched/vtable/kernels.rs b/vortex-array/src/arrays/patched/vtable/kernels.rs index 7994b19e02e..bc05c5e209c 100644 --- a/vortex-array/src/arrays/patched/vtable/kernels.rs +++ b/vortex-array/src/arrays/patched/vtable/kernels.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Patched; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Patched)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Patched)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(Patched), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(Patched)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index 4c0ed1acee7..6d5d70ab6b1 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Canonical; @@ -25,6 +26,8 @@ use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::dtype::DType; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; use crate::scalar_fn::ReduceCtx; @@ -40,10 +43,18 @@ pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[ &ScalarFnAbstractReduceRule, ]); -pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule), - ParentRuleSet::lift(&ScalarFnSliceReduceRule), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 2] = [ + ParentRuleSet::lift_id( + CachedId::new("vortex.filter"), + &ScalarFnUnaryFilterPushDownRule, + ), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &ScalarFnSliceReduceRule), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Converts a ScalarFnArray with Pack into a StructArray directly. #[derive(Debug)] diff --git a/vortex-array/src/arrays/slice/rules.rs b/vortex-array/src/arrays/slice/rules.rs index 3bc0a70f204..29ebebb2c6d 100644 --- a/vortex-array/src/arrays/slice/rules.rs +++ b/vortex-array/src/arrays/slice/rules.rs @@ -1,9 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Slice; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; -pub(super) const PARENT_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(Slice))]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 1] = [ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(Slice), +)]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(super) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 8a802969d39..fe0557e0783 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -4,6 +4,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::IntoArray; @@ -19,6 +20,8 @@ use crate::arrays::slice::SliceReduceAdaptor; use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::cast::Cast; @@ -27,13 +30,18 @@ use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; use crate::validity::Validity; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&StructCastPushDownRule), - ParentRuleSet::lift(&StructGetItemRule), - ParentRuleSet::lift(&MaskReduceAdaptor(Struct)), - ParentRuleSet::lift(&SliceReduceAdaptor(Struct)), - ParentRuleSet::lift(&TakeReduceAdaptor(Struct)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 5] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &StructCastPushDownRule), + ParentRuleSet::lift_id(CachedId::new("vortex.get_item"), &StructGetItemRule), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(Struct)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(Struct)), + ParentRuleSet::lift_id(CachedId::new("vortex.dict"), &TakeReduceAdaptor(Struct)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); /// Rule to push down cast into struct fields. /// diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index 1c5c9f3db3a..f5724da1107 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -1,12 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::Struct; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(Struct)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Struct)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id(CachedId::new("vortex.cast"), &CastExecuteAdaptor(Struct)), + ParentKernelSet::lift_id(CachedId::new("vortex.zip"), &ZipExecuteAdaptor(Struct)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/varbin/compute/rules.rs b/vortex-array/src/arrays/varbin/compute/rules.rs index 89e22ce0d20..9ef4e0c7f44 100644 --- a/vortex-array/src/arrays/varbin/compute/rules.rs +++ b/vortex-array/src/arrays/varbin/compute/rules.rs @@ -1,14 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::VarBin; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(VarBin)), - ParentRuleSet::lift(&MaskReduceAdaptor(VarBin)), - ParentRuleSet::lift(&SliceReduceAdaptor(VarBin)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(VarBin)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(VarBin)), + ParentRuleSet::lift_id(CachedId::new("vortex.slice"), &SliceReduceAdaptor(VarBin)), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 8cb6a2ca377..de74e1dd8e9 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -1,14 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::VarBin; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&FilterExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBin)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 3] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.binary"), + &CompareExecuteAdaptor(VarBin), + ), + ParentKernelSet::lift_id( + CachedId::new("vortex.filter"), + &FilterExecuteAdaptor(VarBin), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.dict"), &TakeExecuteAdaptor(VarBin)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/arrays/varbinview/compute/rules.rs b/vortex-array/src/arrays/varbinview/compute/rules.rs index 5ec24dca7de..d49ba6c69a7 100644 --- a/vortex-array/src/arrays/varbinview/compute/rules.rs +++ b/vortex-array/src/arrays/varbinview/compute/rules.rs @@ -1,13 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_session::registry::CachedId; + use crate::arrays::VarBinView; use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleDense; +use crate::optimizer::rules::ParentRuleEntry; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::cast::CastReduceAdaptor; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&CastReduceAdaptor(VarBinView)), - ParentRuleSet::lift(&MaskReduceAdaptor(VarBinView)), - ParentRuleSet::lift(&SliceReduceAdaptor(VarBinView)), -]); +static KEYED_PARENT_RULES: [ParentRuleEntry; 3] = [ + ParentRuleSet::lift_id(CachedId::new("vortex.cast"), &CastReduceAdaptor(VarBinView)), + ParentRuleSet::lift_id(CachedId::new("vortex.mask"), &MaskReduceAdaptor(VarBinView)), + ParentRuleSet::lift_id( + CachedId::new("vortex.slice"), + &SliceReduceAdaptor(VarBinView), + ), +]; + +static KEYED_PARENT_RULES_DENSE: ParentRuleDense = ParentRuleDense::new(); + +pub(crate) static PARENT_RULES: ParentRuleSet = + ParentRuleSet::new_indexed(&KEYED_PARENT_RULES, &KEYED_PARENT_RULES_DENSE, &[]); diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index 8486f0959f7..04a5733f991 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::registry::CachedId; + use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::kernel::ParentKernelDense; +use crate::kernel::ParentKernelEntry; use crate::kernel::ParentKernelSet; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBinView)), - ParentKernelSet::lift(&ZipExecuteAdaptor(VarBinView)), -]); +static KEYED_PARENT_KERNELS: [ParentKernelEntry; 2] = [ + ParentKernelSet::lift_id( + CachedId::new("vortex.dict"), + &TakeExecuteAdaptor(VarBinView), + ), + ParentKernelSet::lift_id(CachedId::new("vortex.zip"), &ZipExecuteAdaptor(VarBinView)), +]; + +static KEYED_PARENT_KERNELS_DENSE: ParentKernelDense = ParentKernelDense::new(); + +pub(super) static PARENT_KERNELS: ParentKernelSet = + ParentKernelSet::new_indexed(&KEYED_PARENT_KERNELS, &KEYED_PARENT_KERNELS_DENSE, &[]); diff --git a/vortex-array/src/kernel.rs b/vortex-array/src/kernel.rs index 5ad9849ddc4..cd03918d80f 100644 --- a/vortex-array/src/kernel.rs +++ b/vortex-array/src/kernel.rs @@ -40,12 +40,12 @@ pub struct ParentKernelSet { } /// A parent kernel keyed by exact parent encoding id. -pub(crate) struct ParentKernelEntry { +pub struct ParentKernelEntry { parent_id: CachedId, kernel: &'static dyn DynParentKernel, } -pub(crate) type ParentKernelDense = OnceLock>]>>; +pub type ParentKernelDense = OnceLock>]>>; impl ParentKernelSet { fn kernels_for(&self, parent_id: ArrayId) -> &[&'static dyn DynParentKernel] { @@ -81,7 +81,7 @@ impl ParentKernelSet { } } - pub(crate) const fn new_indexed( + pub const fn new_indexed( keyed: &'static [ParentKernelEntry], dense: &'static ParentKernelDense, dynamic: &'static [&'static dyn DynParentKernel], @@ -106,7 +106,7 @@ impl ParentKernelSet { unsafe { &*(kernel as *const K as *const ParentKernelAdapter) } } - pub(crate) const fn lift_id>( + pub const fn lift_id>( parent_id: CachedId, kernel: &'static K, ) -> ParentKernelEntry { diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index 3e5ba173fc5..6a3474e11d0 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -151,12 +151,12 @@ pub struct ParentRuleSet { } /// A parent reduction rule keyed by exact parent encoding id. -pub(crate) struct ParentRuleEntry { +pub struct ParentRuleEntry { parent_id: CachedId, rule: &'static dyn DynArrayParentReduceRule, } -pub(crate) type ParentRuleDense = OnceLock>]>>; +pub type ParentRuleDense = OnceLock>]>>; impl ParentRuleSet { fn rules_for(&self, parent_id: ArrayId) -> &[&'static dyn DynArrayParentReduceRule] { @@ -192,7 +192,7 @@ impl ParentRuleSet { } } - pub(crate) const fn new_indexed( + pub const fn new_indexed( keyed: &'static [ParentRuleEntry], dense: &'static ParentRuleDense, dynamic: &'static [&'static dyn DynArrayParentReduceRule], @@ -217,7 +217,7 @@ impl ParentRuleSet { unsafe { &*(rule as *const R as *const ParentReduceRuleAdapter) } } - pub(crate) const fn lift_id>( + pub const fn lift_id>( parent_id: CachedId, rule: &'static R, ) -> ParentRuleEntry { From 2e02412898f3bc3df33b7a109167b18d4f820389 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Apr 2026 16:14:22 -0400 Subject: [PATCH 09/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/arrays/chunked/vtable/canonical.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index f582054763d..bee94e7e894 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -46,7 +46,7 @@ pub(super) fn _canonicalize( Ok(match array.dtype() { DType::Struct(struct_dtype, _) => { let struct_array = - pack_struct_chunks(&owned_chunks, array.array().validity()?, struct_dtype, ctx)?; + pack_struct_chunks(&owned_chunks, array.array().validity()?, struct_dtype)?; Canonical::Struct(struct_array) } DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( From 2ed04f7185f2764443129e049ca8b854ba442638 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 17:35:37 -0400 Subject: [PATCH 10/16] fix Signed-off-by: Joe Isaacs --- vortex-array/public-api.lock | 38 +- vortex-array/src/arrays/chunked/vtable/mod.rs | 18 +- vortex-array/src/executor.rs | 564 ++++++++++++------ 3 files changed, 421 insertions(+), 199 deletions(-) diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index 88915460d76..d184f781fc9 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -1458,10 +1458,6 @@ pub mod vortex_array::arrays::constant pub struct vortex_array::arrays::constant::Constant -impl vortex_array::arrays::Constant - -pub const vortex_array::arrays::Constant::TAKE_RULES: vortex_array::optimizer::rules::ParentRuleSet - impl core::clone::Clone for vortex_array::arrays::Constant pub fn vortex_array::arrays::Constant::clone(&self) -> vortex_array::arrays::Constant @@ -3216,10 +3212,6 @@ pub mod vortex_array::arrays::null pub struct vortex_array::arrays::null::Null -impl vortex_array::arrays::null::Null - -pub const vortex_array::arrays::null::Null::TAKE_RULES: vortex_array::optimizer::rules::ParentRuleSet - impl core::clone::Clone for vortex_array::arrays::null::Null pub fn vortex_array::arrays::null::Null::clone(&self) -> vortex_array::arrays::null::Null @@ -5276,10 +5268,6 @@ pub fn vortex_array::arrays::Chunked::zip(if_true: vortex_array::ArrayView<'_, v pub struct vortex_array::arrays::Constant -impl vortex_array::arrays::Constant - -pub const vortex_array::arrays::Constant::TAKE_RULES: vortex_array::optimizer::rules::ParentRuleSet - impl core::clone::Clone for vortex_array::arrays::Constant pub fn vortex_array::arrays::Constant::clone(&self) -> vortex_array::arrays::Constant @@ -5988,10 +5976,6 @@ pub fn vortex_array::arrays::Masked::mask(array: vortex_array::ArrayView<'_, vor pub struct vortex_array::arrays::Null -impl vortex_array::arrays::null::Null - -pub const vortex_array::arrays::null::Null::TAKE_RULES: vortex_array::optimizer::rules::ParentRuleSet - impl core::clone::Clone for vortex_array::arrays::null::Null pub fn vortex_array::arrays::null::Null::clone(&self) -> vortex_array::arrays::null::Null @@ -12956,6 +12940,8 @@ pub fn vortex_array::kernel::ParentKernelAdapter::execute_parent(&self, ch pub fn vortex_array::kernel::ParentKernelAdapter::matches(&self, parent: &vortex_array::ArrayRef) -> bool +pub struct vortex_array::kernel::ParentKernelEntry + pub struct vortex_array::kernel::ParentKernelSet impl vortex_array::kernel::ParentKernelSet @@ -12964,8 +12950,12 @@ pub fn vortex_array::kernel::ParentKernelSet::execute(&self, child: vortex_ar pub const fn vortex_array::kernel::ParentKernelSet::lift>(kernel: &'static K) -> &'static dyn vortex_array::kernel::DynParentKernel +pub const fn vortex_array::kernel::ParentKernelSet::lift_id>(parent_id: vortex_session::registry::CachedId, kernel: &'static K) -> vortex_array::kernel::ParentKernelEntry + pub const fn vortex_array::kernel::ParentKernelSet::new(kernels: &'static [&'static dyn vortex_array::kernel::DynParentKernel]) -> Self +pub const fn vortex_array::kernel::ParentKernelSet::new_indexed(keyed: &'static [vortex_array::kernel::ParentKernelEntry], dense: &'static vortex_array::kernel::ParentKernelDense, dynamic: &'static [&'static dyn vortex_array::kernel::DynParentKernel]) -> Self + pub trait vortex_array::kernel::DynParentKernel: core::marker::Send + core::marker::Sync pub fn vortex_array::kernel::DynParentKernel::execute_parent(&self, child: vortex_array::ArrayView<'_, V>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -13056,6 +13046,8 @@ pub type vortex_array::scalar_fn::fns::zip::ZipExecuteAdaptor::Parent = vorte pub fn vortex_array::scalar_fn::fns::zip::ZipExecuteAdaptor::execute_parent(&self, array: vortex_array::ArrayView<'_, V>, parent: vortex_array::arrays::scalar_fn::ScalarFnArrayView<'_, vortex_array::scalar_fn::fns::zip::Zip>, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> +pub type vortex_array::kernel::ParentKernelDense = std::sync::once_lock::OnceLock>]>> + pub mod vortex_array::mask pub mod vortex_array::matcher @@ -13264,6 +13256,8 @@ impl::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct vortex_array::optimizer::rules::ParentRuleEntry + pub struct vortex_array::optimizer::rules::ParentRuleSet impl vortex_array::optimizer::rules::ParentRuleSet @@ -13272,8 +13266,12 @@ pub fn vortex_array::optimizer::rules::ParentRuleSet::evaluate(&self, child: pub const fn vortex_array::optimizer::rules::ParentRuleSet::lift>(rule: &'static R) -> &'static dyn vortex_array::optimizer::rules::DynArrayParentReduceRule +pub const fn vortex_array::optimizer::rules::ParentRuleSet::lift_id>(parent_id: vortex_session::registry::CachedId, rule: &'static R) -> vortex_array::optimizer::rules::ParentRuleEntry + pub const fn vortex_array::optimizer::rules::ParentRuleSet::new(rules: &'static [&'static dyn vortex_array::optimizer::rules::DynArrayParentReduceRule]) -> Self +pub const fn vortex_array::optimizer::rules::ParentRuleSet::new_indexed(keyed: &'static [vortex_array::optimizer::rules::ParentRuleEntry], dense: &'static vortex_array::optimizer::rules::ParentRuleDense, dynamic: &'static [&'static dyn vortex_array::optimizer::rules::DynArrayParentReduceRule]) -> Self + pub struct vortex_array::optimizer::rules::ReduceRuleSet impl vortex_array::optimizer::rules::ReduceRuleSet @@ -13388,6 +13386,8 @@ pub fn vortex_array::optimizer::rules::ParentReduceRuleAdapter::matches(&s pub fn vortex_array::optimizer::rules::ParentReduceRuleAdapter::reduce_parent(&self, child: vortex_array::ArrayView<'_, V>, parent: &vortex_array::ArrayRef, child_idx: usize) -> vortex_error::VortexResult> +pub type vortex_array::optimizer::rules::ParentRuleDense = std::sync::once_lock::OnceLock>]>> + pub trait vortex_array::optimizer::ArrayOptimizer pub fn vortex_array::optimizer::ArrayOptimizer::optimize(&self) -> vortex_error::VortexResult @@ -21438,6 +21438,8 @@ pub vortex_array::ExecutionStep::Done pub vortex_array::ExecutionStep::ExecuteSlot(usize, vortex_array::DonePredicate) +pub vortex_array::ExecutionStep::ExecuteSlots(alloc::vec::Vec<(usize, vortex_array::DonePredicate)>) + impl core::fmt::Debug for vortex_array::ExecutionStep pub fn vortex_array::ExecutionStep::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -22552,8 +22554,12 @@ pub fn vortex_array::ExecutionResult::array(&self) -> &vortex_array::ArrayRef pub fn vortex_array::ExecutionResult::done(result: impl vortex_array::IntoArray) -> Self +pub fn vortex_array::ExecutionResult::execute_range(array: impl vortex_array::IntoArray, slots: core::ops::range::Range) -> Self + pub fn vortex_array::ExecutionResult::execute_slot(array: impl vortex_array::IntoArray, slot_idx: usize) -> Self +pub fn vortex_array::ExecutionResult::execute_slots(array: impl vortex_array::IntoArray, slots: impl core::iter::traits::collect::IntoIterator) -> Self + pub fn vortex_array::ExecutionResult::into_parts(self) -> (vortex_array::ArrayRef, vortex_array::ExecutionStep) pub fn vortex_array::ExecutionResult::step(&self) -> &vortex_array::ExecutionStep diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 1171f40d230..f42a96028f6 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -240,16 +240,14 @@ impl VTable for Chunked { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - // Iteratively request execution of each chunk until it reaches canonical form. - // This gives the scheduler visibility into child execution and enables - // cross-step optimizations between chunk decoding steps. - for i in 0..array.nchunks() { - if !array.chunk(i).is::() { - return Ok(ExecutionResult::execute_slot::( - array, - i + CHUNKS_OFFSET, - )); - } + // Request all chunks together so the scheduler can keep this parent frame on the stack + // while it walks the homogeneous chunk slots. + if (0..array.nchunks()).any(|i| !array.chunk(i).is::()) { + let chunk_slots = CHUNKS_OFFSET..CHUNKS_OFFSET + array.nchunks(); + return Ok(ExecutionResult::execute_range::( + array, + chunk_slots, + )); } // All chunks are now canonical — combine them. diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 3f48c01f7c7..adb800e2d85 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -20,6 +20,7 @@ use std::env::VarError; use std::fmt; use std::fmt::Display; +use std::ops::Range; use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; @@ -86,83 +87,92 @@ impl ArrayRef { /// Iteratively execute this array until the [`Matcher`] matches, using an explicit work /// stack. /// - /// Each iteration proceeds through three steps in order: + /// The stack is a flat `Vec` processed back-to-front. Each entry holds an array + /// and optionally links back to a parent entry (by index) where the result should be + /// placed when done. Entries without a parent link are the root. /// - /// 1. **Done / canonical check** — if `current` satisfies the active done predicate or is - /// canonical, splice it back into the stacked parent (if any) and continue, or return. - /// 2. **`execute_parent` on children** — try each child's `execute_parent` against `current` - /// as the parent (e.g. `Filter(RunEnd)` → `FilterExecuteAdaptor` fires from RunEnd). - /// If there is a stacked parent frame, the rewritten child is spliced back into it so - /// that optimize and further `execute_parent` can fire on the reconstructed parent - /// (e.g. `Slice(RunEnd)` → `RunEnd` spliced into stacked `Filter` → `Filter(RunEnd)` - /// whose `FilterExecuteAdaptor` fires on the next iteration). - /// 3. **`execute`** — call the encoding's own execute step, which either returns `Done` or - /// `ExecuteSlot(i)` to push a child onto the stack for focused execution. - /// - /// Note: the returned array may not match `M`. If execution converges to a canonical form - /// that does not match `M`, the canonical array is returned since no further execution - /// progress is possible. + /// When an encoding returns `ExecuteSlot` or `ExecuteSlots`, the current array is pushed + /// back as a parent entry, its children are extracted and pushed in reverse order (so the + /// first child is on top and processed first). Parent indices are stable because entries + /// are only pushed/popped at the top. /// /// For safety, we will error when the number of execution iterations reaches a configurable /// maximum (default 128, override with `VORTEX_MAX_ITERATIONS`). pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { - let mut current = self.optimize()?; - let mut stack: Vec = Vec::new(); + let root = self.optimize()?; + let mut stack: Vec = vec![Entry::root(root, M::matches)]; for _ in 0..max_iterations() { - // Step 1: done / canonical — splice back into stacked parent or return. - let is_done = stack - .last() - .map_or(M::matches as DonePredicate, |frame| frame.done); - if is_done(¤t) || AnyCanonical::matches(¤t) { - match stack.pop() { + let mut entry = match stack.pop() { + Some(e) => e, + None => vortex_bail!("Empty execution stack"), + }; + let current = entry.array.take().expect("entry must have array"); + + // Step 1: done / canonical — put back into parent or return. + if (entry.done)(¤t) || AnyCanonical::matches(¤t) { + match entry.parent { None => { ctx.log(format_args!("-> {}", current)); return Ok(current); } - Some(frame) => { - current = frame.put_back(current)?.optimize()?; + Some(ref link) => { + put_back_child( + &mut stack, + link.parent_idx, + link.slot_idx, + current, + &entry, + )?; continue; } } } - // Step 2: execute_parent on children (current is the parent). - // If there is a stacked parent frame, splice the rewritten child back into it - // so that optimize and execute_parent can fire naturally on the reconstructed parent - // (e.g. Slice(RunEnd) -RunEndSliceKernel-> RunEnd, spliced back into Filter gives - // Filter(RunEnd), whose FilterExecuteAdaptor fires on the next iteration). - if let Some(rewritten) = try_execute_parent(¤t, ctx)? { - ctx.log(format_args!( - "execute_parent rewrote {} -> {}", - current, rewritten - )); - current = rewritten.optimize()?; - if let Some(frame) = stack.pop() { - current = frame.put_back(current)?.optimize()?; - } - continue; - } - - // Step 4: execute the encoding's own step. - let result = execute_step(current, ctx)?; - let (array, step) = result.into_parts(); - match step { - ExecutionStep::ExecuteSlot(i, done) => { - // SAFETY: we record the child's dtype and len, and assert they are preserved - // when the slot is put back via `put_slot_unchecked`. - let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "ExecuteSlot({i}): pushing {}, focusing on {}", - parent, child - )); - let frame = StackFrame::new(parent, i, done, &child); - stack.push(frame); - current = child.optimize()?; + // Step 2: execute_parent, then execute encoding's own step. + let result = execute_one_step(current, ctx)?; + match result { + OneStepResult::Rewritten(rewritten) => { + let rewritten = rewritten.optimize()?; + match entry.parent { + None => { + entry.array = Some(rewritten); + stack.push(entry); + } + Some(link) => { + unwind_to_parent(&mut stack, link, rewritten, ctx)?; + } + } } - ExecutionStep::Done => { - ctx.log(format_args!("Done: {}", array)); - current = array; + OneStepResult::Execute(result) => { + let (array, step) = result.into_parts(); + match step { + ExecutionStep::ExecuteSlot(i, done) => { + push_children( + array, + vec![(i, done)], + entry.parent, + entry.done, + &mut stack, + ctx, + )?; + } + ExecutionStep::ExecuteSlots(slots) => { + push_children( + array, + slots, + entry.parent, + entry.done, + &mut stack, + ctx, + )?; + } + ExecutionStep::Done => { + ctx.log(format_args!("Done: {}", array)); + entry.array = Some(array); + stack.push(entry); + } + } } } } @@ -174,49 +184,199 @@ impl ArrayRef { } } -/// A stack frame for the iterative executor, tracking the parent array whose slot is being -/// executed and the original child's dtype/len for validation on put-back. -struct StackFrame { - parent: ArrayRef, - slot_idx: usize, +/// A work item on the execution stack. +/// +/// Each entry holds an array and optionally links back to a parent entry (by index) where +/// the result should be placed when done. Entries without a parent link are the root. +struct Entry { + array: Option, + parent: Option, done: DonePredicate, original_dtype: DType, original_len: usize, } -impl StackFrame { - fn new(parent: ArrayRef, slot_idx: usize, done: DonePredicate, child: &ArrayRef) -> Self { +/// Points from a child [`Entry`] back to its parent in the stack. +struct ParentLink { + /// Index of the parent [`Entry`] in the stack vec. + parent_idx: usize, + /// Which slot in the parent array to put this child's result into. + slot_idx: usize, +} + +impl Entry { + fn root(array: ArrayRef, done: DonePredicate) -> Self { + let dtype = array.dtype().clone(); + let len = array.len(); Self { - parent, - slot_idx, + array: Some(array), + parent: None, done, - original_dtype: child.dtype().clone(), - original_len: child.len(), + original_dtype: dtype, + original_len: len, } } - fn put_back(self, replacement: ArrayRef) -> VortexResult { - debug_assert_eq!( - replacement.dtype(), - &self.original_dtype, - "slot {} dtype changed from {} to {} during execution", - self.slot_idx, - self.original_dtype, - replacement.dtype() - ); - debug_assert_eq!( - replacement.len(), - self.original_len, - "slot {} len changed from {} to {} during execution", - self.slot_idx, - self.original_len, - replacement.len() - ); - // SAFETY: we assert above that dtype and len are preserved. - unsafe { self.parent.put_slot_unchecked(self.slot_idx, replacement) } + fn child( + array: ArrayRef, + parent_idx: usize, + slot_idx: usize, + done: DonePredicate, + ) -> Self { + let dtype = array.dtype().clone(); + let len = array.len(); + Self { + array: Some(array), + parent: Some(ParentLink { + parent_idx, + slot_idx, + }), + done, + original_dtype: dtype, + original_len: len, + } } } +/// Put a completed child back into its parent entry in the stack. +fn put_back_child( + stack: &mut [Entry], + parent_idx: usize, + slot_idx: usize, + child: ArrayRef, + child_entry: &Entry, +) -> VortexResult<()> { + debug_assert_eq!( + child.dtype(), + &child_entry.original_dtype, + "slot {} dtype changed from {} to {} during execution", + slot_idx, + child_entry.original_dtype, + child.dtype() + ); + debug_assert_eq!( + child.len(), + child_entry.original_len, + "slot {} len changed from {} to {} during execution", + slot_idx, + child_entry.original_len, + child.len() + ); + let parent_array = stack[parent_idx] + .array + .take() + .expect("parent must have array"); + // SAFETY: we assert above that dtype and len are preserved. + let updated = unsafe { parent_array.put_slot_unchecked(slot_idx, child)? }; + stack[parent_idx].array = Some(updated); + Ok(()) +} + +/// Unwind: put all entries above the parent back into their respective parents, +/// then put the rewritten child back. This collapses any partially-expanded +/// descendants so the parent can be re-processed with all children in place. +fn unwind_to_parent( + stack: &mut Vec, + link: ParentLink, + rewritten_child: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + // Pop entries above the parent, putting each back into its own parent. + // Deepest entries are on top so grandchildren are put back before siblings. + while stack.len() > link.parent_idx + 1 { + let entry = stack.pop().unwrap(); + if let (Some(entry_link), Some(entry_array)) = (entry.parent, entry.array) { + // SAFETY: these entries haven't been modified, dtype/len are preserved. + let parent_array = stack[entry_link.parent_idx] + .array + .take() + .expect("parent must have array"); + let updated = + unsafe { parent_array.put_slot_unchecked(entry_link.slot_idx, entry_array)? }; + stack[entry_link.parent_idx].array = Some(updated); + } + } + + // Put the rewritten child back into the parent. + let parent_array = stack[link.parent_idx] + .array + .take() + .expect("parent must have array"); + let updated = unsafe { parent_array.put_slot_unchecked(link.slot_idx, rewritten_child)? }; + let updated = updated.optimize()?; + ctx.log(format_args!("unwind -> {}", updated)); + stack[link.parent_idx].array = Some(updated); + Ok(()) +} + +/// Extract children from a parent array and push them onto the stack. +/// +/// The parent is pushed first (at a stable index), then children are pushed in +/// reverse order so the first slot to process is on top. +fn push_children( + mut parent: ArrayRef, + slots: Vec<(usize, DonePredicate)>, + parent_parent: Option, + parent_done: DonePredicate, + stack: &mut Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + // Collect children that actually need execution. + let mut to_push: Vec<(usize, DonePredicate, ArrayRef)> = Vec::new(); + for (slot_idx, done) in &slots { + let Some(child) = parent.slots().get(*slot_idx).and_then(Option::as_ref) else { + vortex_bail!( + "Execution requested slot {} but array {} has no occupied slot there", + slot_idx, + parent + ); + }; + if (done)(child) || AnyCanonical::matches(child) { + continue; + } + let (new_parent, child) = unsafe { parent.take_slot_unchecked(*slot_idx) }?; + parent = new_parent; + to_push.push((*slot_idx, *done, child)); + } + + if to_push.is_empty() { + // All children already done, push parent back as-is. + let dtype = parent.dtype().clone(); + let len = parent.len(); + stack.push(Entry { + array: Some(parent), + parent: parent_parent, + done: parent_done, + original_dtype: dtype, + original_len: len, + }); + return Ok(()); + } + + // Push parent entry at a stable index. + let parent_idx = stack.len(); + let dtype = parent.dtype().clone(); + let len = parent.len(); + stack.push(Entry { + array: Some(parent), + parent: parent_parent, + done: parent_done, + original_dtype: dtype, + original_len: len, + }); + + // Push children in reverse order (first-to-process on top). + for (slot_idx, done, child) in to_push.into_iter().rev() { + let child = child.optimize()?; + ctx.log(format_args!( + "ExecuteSlot({slot_idx}): pushing, focusing on {child}" + )); + stack.push(Entry::child(child, parent_idx, slot_idx, done)); + } + + Ok(()) +} + /// Execution context for batch CPU compute. #[derive(Debug, Clone)] pub struct ExecutionCtx { @@ -295,108 +455,125 @@ impl Drop for ExecutionCtx { /// is incrementally more "executed" than the input array. In other words, it is closer to becoming /// a canonical array. /// -/// The execution steps are as follows: -/// 0. Check for canonical. -/// 1. Attempt to `reduce` the array with metadata-only optimizations. -/// 2. Attempt to call `reduce_parent` on each child. -/// 3. Attempt to call `execute_parent` on each child. -/// 4. Call `execute` on the array itself (which returns an [`ExecutionStep`]). -/// /// Most users will not call this method directly, instead preferring to specify an executable /// target such as [`crate::Columnar`], [`Canonical`], or any of the canonical array types (such as /// [`crate::arrays::PrimitiveArray`]). impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // 0. Check for canonical - if let Some(canonical) = array.as_opt::() { - ctx.log(format_args!("-> canonical {}", array)); - return Ok(Canonical::from(canonical).into_array()); + let result = execute_one_step(array, ctx)?; + match result { + OneStepResult::Rewritten(array) => Ok(array), + OneStepResult::Execute(result) => { + let (array, step) = result.into_parts(); + match step { + ExecutionStep::Done => { + ctx.log(format_args!("-> {}", array)); + Ok(array) + } + ExecutionStep::ExecuteSlot(i, _) => { + let (array, child) = unsafe { array.take_slot_unchecked(i)? }; + let executed = child.execute::(ctx)?; + unsafe { array.put_slot_unchecked(i, executed) } + } + ExecutionStep::ExecuteSlots(slots) => { + // Single-step: execute only the first undone child. + for (slot_idx, done) in slots { + let already_done = array + .slots() + .get(slot_idx) + .and_then(Option::as_ref) + .is_some_and(|c| done(c) || AnyCanonical::matches(c)); + if already_done { + continue; + } + let (parent, child) = + unsafe { array.take_slot_unchecked(slot_idx)? }; + let executed = child.execute::(ctx)?; + return unsafe { parent.put_slot_unchecked(slot_idx, executed) }; + } + Ok(array) + } + } + } } + } +} - // 1. reduce (metadata-only rewrites) - if let Some(reduced) = array.reduce()? { - ctx.log(format_args!("reduce: rewrote {} -> {}", array, reduced)); - reduced.statistics().inherit_from(array.statistics()); - return Ok(reduced); - } +/// The result of a single execution step, before the caller decides how to handle it. +enum OneStepResult { + /// The array was rewritten by a canonical check, reduce, reduce_parent, or execute_parent step. + Rewritten(ArrayRef), + /// The encoding's own `execute` step ran and returned an [`ExecutionResult`]. + Execute(ExecutionResult), +} - // 2. reduce_parent (child-driven metadata-only rewrites) - for (slot_idx, slot) in array.slots().iter().enumerate() { - let Some(child) = slot else { - continue; - }; - if let Some(reduced_parent) = child.reduce_parent(&array, slot_idx)? { - ctx.log(format_args!( - "reduce_parent: slot[{}]({}) rewrote {} -> {}", - slot_idx, - child.encoding_id(), - array, - reduced_parent - )); - reduced_parent.statistics().inherit_from(array.statistics()); - return Ok(reduced_parent); - } - } +/// Perform one step of execution on an array, trying each layer in priority order: +/// +/// 0. Check for canonical. +/// 1. `reduce` — metadata-only self-rewrite. +/// 2. `reduce_parent` — metadata-only child-driven parent rewrite. +/// 3. `execute_parent` — child-driven fused execution (may read buffers). +/// 4. `execute` — the encoding's own decode step. +/// +/// Returns a [`OneStepResult`] so the caller can decide how to handle the outcome +/// (single-step vs iterative loop with stack). +fn execute_one_step(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // 0. Check for canonical + if let Some(canonical) = array.as_opt::() { + ctx.log(format_args!("-> canonical {}", array)); + return Ok(OneStepResult::Rewritten( + Canonical::from(canonical).into_array(), + )); + } - // 3. execute_parent (child-driven optimized execution) - for (slot_idx, slot) in array.slots().iter().enumerate() { - let Some(child) = slot else { - continue; - }; - if let Some(executed_parent) = child.execute_parent(&array, slot_idx, ctx)? { - ctx.log(format_args!( - "execute_parent: slot[{}]({}) rewrote {} -> {}", - slot_idx, - child.encoding_id(), - array, - executed_parent - )); - executed_parent - .statistics() - .inherit_from(array.statistics()); - return Ok(executed_parent); - } - } + // 1. reduce (metadata-only rewrites) + if let Some(reduced) = array.reduce()? { + ctx.log(format_args!("reduce: rewrote {} -> {}", array, reduced)); + reduced.statistics().inherit_from(array.statistics()); + return Ok(OneStepResult::Rewritten(reduced)); + } - // 4. execute (returns an ExecutionResult) - ctx.log(format_args!("executing {}", array)); - let result = execute_step(array, ctx)?; - let (array, step) = result.into_parts(); - match step { - ExecutionStep::Done => { - ctx.log(format_args!("-> {}", array)); - Ok(array) - } - ExecutionStep::ExecuteSlot(i, _) => { - // For single-step execution, handle ExecuteSlot by executing the slot, - // replacing it, and returning the updated array. - let (array, child) = unsafe { array.take_slot_unchecked(i)? }; - let executed_child = child.execute::(ctx)?; - unsafe { array.put_slot_unchecked(i, executed_child) } - } + // 2. reduce_parent (child-driven metadata-only rewrites) + for (slot_idx, slot) in array.slots().iter().enumerate() { + let Some(child) = slot else { + continue; + }; + if let Some(reduced_parent) = child.reduce_parent(&array, slot_idx)? { + ctx.log(format_args!( + "reduce_parent: slot[{}]({}) rewrote {} -> {}", + slot_idx, + child.encoding_id(), + array, + reduced_parent + )); + reduced_parent.statistics().inherit_from(array.statistics()); + return Ok(OneStepResult::Rewritten(reduced_parent)); } } -} -/// Execute a single step on an array, consuming it. -/// -/// Extracts the vtable before consuming the array to avoid borrow conflicts. -fn execute_step(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - array.execute_encoding(ctx) -} - -/// Try execute_parent on each occupied slot of the array. -fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + // 3. execute_parent (child-driven optimized execution) for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue; }; - if let Some(result) = child.execute_parent(array, slot_idx, ctx)? { - result.statistics().inherit_from(array.statistics()); - return Ok(Some(result)); + if let Some(executed_parent) = child.execute_parent(&array, slot_idx, ctx)? { + ctx.log(format_args!( + "execute_parent: slot[{}]({}) rewrote {} -> {}", + slot_idx, + child.encoding_id(), + array, + executed_parent + )); + executed_parent + .statistics() + .inherit_from(array.statistics()); + return Ok(OneStepResult::Rewritten(executed_parent)); } } - Ok(None) + + // 4. execute (returns an ExecutionResult) + ctx.log(format_args!("executing {}", array)); + Ok(OneStepResult::Execute(array.execute_encoding(ctx)?)) } /// A predicate that determines when an array has reached a desired form during execution. @@ -418,6 +595,17 @@ pub enum ExecutionStep { /// Use [`ExecutionResult::execute_slot`] instead of constructing this variant directly. ExecuteSlot(usize, DonePredicate), + /// Request that the scheduler execute multiple slots, each using its paired + /// [`DonePredicate`], then replace them in this array and re-enter execution. + /// + /// Slots are executed in the order they appear in the vector. The scheduler keeps the parent + /// shape stable until the requested slots are exhausted, because queued slot indices are only + /// meaningful for the parent that produced them. + /// + /// Use [`ExecutionResult::execute_slots`] or [`ExecutionResult::execute_range`] instead of + /// constructing this variant directly. + ExecuteSlots(Vec<(usize, DonePredicate)>), + /// Execution is complete. The array in the accompanying [`ExecutionResult`] is the result. /// The scheduler will continue executing if it has not yet reached the target form. Done, @@ -427,6 +615,10 @@ impl fmt::Debug for ExecutionStep { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { ExecutionStep::ExecuteSlot(idx, _) => f.debug_tuple("ExecuteSlot").field(idx).finish(), + ExecutionStep::ExecuteSlots(slots) => f + .debug_tuple("ExecuteSlots") + .field(&slots.iter().map(|(idx, _)| *idx).collect::>()) + .finish(), ExecutionStep::Done => write!(f, "Done"), } } @@ -460,6 +652,32 @@ impl ExecutionResult { } } + /// Request execution of multiple slots until each matches the given [`Matcher`]. + /// + /// The slots are executed in iterator order. This is useful for encodings whose next + /// execution step requires a batch of homogeneous children, such as all chunks in a chunked + /// array. The iterator must yield at least one slot. + pub fn execute_slots( + array: impl IntoArray, + slots: impl IntoIterator, + ) -> Self { + Self { + array: array.into_array(), + step: ExecutionStep::ExecuteSlots( + slots + .into_iter() + .map(|slot_idx| (slot_idx, M::matches as DonePredicate)) + .collect(), + ), + } + } + + /// Request execution of a non-empty contiguous range of slots until each matches the given + /// [`Matcher`]. + pub fn execute_range(array: impl IntoArray, slots: Range) -> Self { + Self::execute_slots::(array, slots) + } + /// Returns a reference to the array. pub fn array(&self) -> &ArrayRef { &self.array From 346bd719deb8c9b4792bf5654d98e92e8a1ee2dc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 17:54:51 -0400 Subject: [PATCH 11/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 122 +++++++++-------------------------- 1 file changed, 32 insertions(+), 90 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index adb800e2d85..2155ff5213b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -140,7 +140,7 @@ impl ArrayRef { stack.push(entry); } Some(link) => { - unwind_to_parent(&mut stack, link, rewritten, ctx)?; + put_back_rewritten(&mut stack, link, rewritten)?; } } } @@ -158,14 +158,7 @@ impl ArrayRef { )?; } ExecutionStep::ExecuteSlots(slots) => { - push_children( - array, - slots, - entry.parent, - entry.done, - &mut stack, - ctx, - )?; + push_children(array, slots, entry.parent, entry.done, &mut stack, ctx)?; } ExecutionStep::Done => { ctx.log(format_args!("Done: {}", array)); @@ -197,6 +190,7 @@ struct Entry { } /// Points from a child [`Entry`] back to its parent in the stack. +#[derive(Clone)] struct ParentLink { /// Index of the parent [`Entry`] in the stack vec. parent_idx: usize, @@ -217,12 +211,7 @@ impl Entry { } } - fn child( - array: ArrayRef, - parent_idx: usize, - slot_idx: usize, - done: DonePredicate, - ) -> Self { + fn child(array: ArrayRef, parent_idx: usize, slot_idx: usize, done: DonePredicate) -> Self { let dtype = array.dtype().clone(); let len = array.len(); Self { @@ -272,108 +261,62 @@ fn put_back_child( Ok(()) } -/// Unwind: put all entries above the parent back into their respective parents, -/// then put the rewritten child back. This collapses any partially-expanded -/// descendants so the parent can be re-processed with all children in place. -fn unwind_to_parent( - stack: &mut Vec, - link: ParentLink, - rewritten_child: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult<()> { - // Pop entries above the parent, putting each back into its own parent. - // Deepest entries are on top so grandchildren are put back before siblings. - while stack.len() > link.parent_idx + 1 { - let entry = stack.pop().unwrap(); - if let (Some(entry_link), Some(entry_array)) = (entry.parent, entry.array) { - // SAFETY: these entries haven't been modified, dtype/len are preserved. - let parent_array = stack[entry_link.parent_idx] - .array - .take() - .expect("parent must have array"); - let updated = - unsafe { parent_array.put_slot_unchecked(entry_link.slot_idx, entry_array)? }; - stack[entry_link.parent_idx].array = Some(updated); - } - } - - // Put the rewritten child back into the parent. - let parent_array = stack[link.parent_idx] +/// Put a rewritten child back into its parent. Siblings remain on the stack +/// and will put themselves back when they complete. +fn put_back_rewritten(stack: &mut [Entry], link: ParentLink, child: ArrayRef) -> VortexResult<()> { + let parent = stack[link.parent_idx] .array .take() .expect("parent must have array"); - let updated = unsafe { parent_array.put_slot_unchecked(link.slot_idx, rewritten_child)? }; - let updated = updated.optimize()?; - ctx.log(format_args!("unwind -> {}", updated)); + let updated = unsafe { parent.put_slot_unchecked(link.slot_idx, child)? }; stack[link.parent_idx].array = Some(updated); Ok(()) } -/// Extract children from a parent array and push them onto the stack. -/// -/// The parent is pushed first (at a stable index), then children are pushed in -/// reverse order so the first slot to process is on top. +/// Extract all undone children from the parent and push them onto the stack in +/// reverse order (first-to-process on top). Each child puts itself back into the +/// parent via [`put_back_child`] when done. fn push_children( mut parent: ArrayRef, - slots: Vec<(usize, DonePredicate)>, + mut slots: Vec<(usize, DonePredicate)>, parent_parent: Option, parent_done: DonePredicate, stack: &mut Vec, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - // Collect children that actually need execution. - let mut to_push: Vec<(usize, DonePredicate, ArrayRef)> = Vec::new(); - for (slot_idx, done) in &slots { - let Some(child) = parent.slots().get(*slot_idx).and_then(Option::as_ref) else { + let parent_idx = stack.len(); + + // Push parent entry up front; children will reference it by index. + let mut parent_entry = Entry::root(parent, parent_done); + parent_entry.parent = parent_parent; + stack.push(parent_entry); + + // Iterate in reverse so first-to-process ends up on top of the stack. + // Extract each undone child from the parent as we go. + let mut parent = stack[parent_idx].array.take().unwrap(); + slots.reverse(); + for (slot_idx, done) in slots { + let Some(child) = parent.slots().get(slot_idx).and_then(Option::as_ref) else { vortex_bail!( "Execution requested slot {} but array {} has no occupied slot there", slot_idx, parent ); }; - if (done)(child) || AnyCanonical::matches(child) { + if done(child) || AnyCanonical::matches(child) { continue; } - let (new_parent, child) = unsafe { parent.take_slot_unchecked(*slot_idx) }?; + let (new_parent, child) = unsafe { parent.take_slot_unchecked(slot_idx) }?; parent = new_parent; - to_push.push((*slot_idx, *done, child)); - } - - if to_push.is_empty() { - // All children already done, push parent back as-is. - let dtype = parent.dtype().clone(); - let len = parent.len(); - stack.push(Entry { - array: Some(parent), - parent: parent_parent, - done: parent_done, - original_dtype: dtype, - original_len: len, - }); - return Ok(()); - } - - // Push parent entry at a stable index. - let parent_idx = stack.len(); - let dtype = parent.dtype().clone(); - let len = parent.len(); - stack.push(Entry { - array: Some(parent), - parent: parent_parent, - done: parent_done, - original_dtype: dtype, - original_len: len, - }); - - // Push children in reverse order (first-to-process on top). - for (slot_idx, done, child) in to_push.into_iter().rev() { - let child = child.optimize()?; ctx.log(format_args!( "ExecuteSlot({slot_idx}): pushing, focusing on {child}" )); stack.push(Entry::child(child, parent_idx, slot_idx, done)); } + // Store parent back (with extracted slots emptied). + stack[parent_idx].array = Some(parent); + Ok(()) } @@ -486,8 +429,7 @@ impl Executable for ArrayRef { if already_done { continue; } - let (parent, child) = - unsafe { array.take_slot_unchecked(slot_idx)? }; + let (parent, child) = unsafe { array.take_slot_unchecked(slot_idx)? }; let executed = child.execute::(ctx)?; return unsafe { parent.put_slot_unchecked(slot_idx, executed) }; } From 096efa30efa26d1bb89c11a6527c821c3dfe1746 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 17:57:39 -0400 Subject: [PATCH 12/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 2155ff5213b..13641e9b1f5 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -99,8 +99,7 @@ impl ArrayRef { /// For safety, we will error when the number of execution iterations reaches a configurable /// maximum (default 128, override with `VORTEX_MAX_ITERATIONS`). pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { - let root = self.optimize()?; - let mut stack: Vec = vec![Entry::root(root, M::matches)]; + let mut stack: Vec = vec![Entry::root(self, M::matches)]; for _ in 0..max_iterations() { let mut entry = match stack.pop() { From f265d58438713238319efbb24db1ed43d6b64f35 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 18:14:08 -0400 Subject: [PATCH 13/16] fix Signed-off-by: Joe Isaacs --- .../src/arrays/primitive/compute/take/avx2.rs | 1 + vortex-array/src/executor.rs | 55 ++++++------------- 2 files changed, 18 insertions(+), 38 deletions(-) diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs index e92304dc34b..bce14b79f14 100644 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ b/vortex-array/src/arrays/primitive/compute/take/avx2.rs @@ -121,6 +121,7 @@ where /// The caller must ensure the `avx2` feature is enabled. #[target_feature(enable = "avx2")] #[doc(hidden)] +#[inline] unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> Buffer { macro_rules! dispatch_avx2 { ($indices:ty, $values:ty) => { diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 13641e9b1f5..31f52d40062 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -37,7 +37,6 @@ use crate::dtype::DType; use crate::matcher::Matcher; use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; -use crate::optimizer::ArrayOptimizer; /// Returns the maximum number of iterations to attempt when executing an array before giving up and returning /// an error, can be by the `VORTEX_MAX_ITERATIONS` env variables, otherwise defaults to 128. @@ -132,16 +131,8 @@ impl ArrayRef { let result = execute_one_step(current, ctx)?; match result { OneStepResult::Rewritten(rewritten) => { - let rewritten = rewritten.optimize()?; - match entry.parent { - None => { - entry.array = Some(rewritten); - stack.push(entry); - } - Some(link) => { - put_back_rewritten(&mut stack, link, rewritten)?; - } - } + entry.array = Some(rewritten); + stack.push(entry); } OneStepResult::Execute(result) => { let (array, step) = result.into_parts(); @@ -260,61 +251,49 @@ fn put_back_child( Ok(()) } -/// Put a rewritten child back into its parent. Siblings remain on the stack -/// and will put themselves back when they complete. -fn put_back_rewritten(stack: &mut [Entry], link: ParentLink, child: ArrayRef) -> VortexResult<()> { - let parent = stack[link.parent_idx] - .array - .take() - .expect("parent must have array"); - let updated = unsafe { parent.put_slot_unchecked(link.slot_idx, child)? }; - stack[link.parent_idx].array = Some(updated); - Ok(()) -} - /// Extract all undone children from the parent and push them onto the stack in /// reverse order (first-to-process on top). Each child puts itself back into the /// parent via [`put_back_child`] when done. fn push_children( - mut parent: ArrayRef, + child: ArrayRef, mut slots: Vec<(usize, DonePredicate)>, - parent_parent: Option, - parent_done: DonePredicate, + parent: Option, + done: DonePredicate, stack: &mut Vec, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let parent_idx = stack.len(); // Push parent entry up front; children will reference it by index. - let mut parent_entry = Entry::root(parent, parent_done); - parent_entry.parent = parent_parent; + let mut parent_entry = Entry::root(child, done); + parent_entry.parent = parent; stack.push(parent_entry); // Iterate in reverse so first-to-process ends up on top of the stack. // Extract each undone child from the parent as we go. - let mut parent = stack[parent_idx].array.take().unwrap(); + let mut array = stack[parent_idx].array.take().unwrap(); slots.reverse(); - for (slot_idx, done) in slots { - let Some(child) = parent.slots().get(slot_idx).and_then(Option::as_ref) else { + for (slot_idx, slot_done) in slots { + let Some(slot_child) = array.slots().get(slot_idx).and_then(Option::as_ref) else { vortex_bail!( "Execution requested slot {} but array {} has no occupied slot there", slot_idx, - parent + array ); }; - if done(child) || AnyCanonical::matches(child) { + if slot_done(slot_child) || AnyCanonical::matches(slot_child) { continue; } - let (new_parent, child) = unsafe { parent.take_slot_unchecked(slot_idx) }?; - parent = new_parent; + let (new_array, slot_child) = unsafe { array.take_slot_unchecked(slot_idx) }?; + array = new_array; ctx.log(format_args!( - "ExecuteSlot({slot_idx}): pushing, focusing on {child}" + "ExecuteSlot({slot_idx}): pushing, focusing on {slot_child}" )); - stack.push(Entry::child(child, parent_idx, slot_idx, done)); + stack.push(Entry::child(slot_child, parent_idx, slot_idx, slot_done)); } // Store parent back (with extracted slots emptied). - stack[parent_idx].array = Some(parent); + stack[parent_idx].array = Some(array); Ok(()) } From 1a2a8c16529829675762b2ec4f38c4571367424a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 18:21:16 -0400 Subject: [PATCH 14/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 31f52d40062..22110d4012e 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -143,12 +143,15 @@ impl ArrayRef { vec![(i, done)], entry.parent, entry.done, + true, &mut stack, ctx, )?; } ExecutionStep::ExecuteSlots(slots) => { - push_children(array, slots, entry.parent, entry.done, &mut stack, ctx)?; + push_children( + array, slots, entry.parent, entry.done, false, &mut stack, ctx, + )?; } ExecutionStep::Done => { ctx.log(format_args!("Done: {}", array)); @@ -177,6 +180,10 @@ struct Entry { done: DonePredicate, original_dtype: DType, original_len: usize, + /// When true, after this child is put back into its parent, the parent is + /// popped and re-entered into the execute loop so that `execute_parent` can + /// fire on the reconstructed parent. Set for `ExecuteSlot` children only. + reexecute_parent: bool, } /// Points from a child [`Entry`] back to its parent in the stack. @@ -198,6 +205,7 @@ impl Entry { done, original_dtype: dtype, original_len: len, + reexecute_parent: false, } } @@ -213,6 +221,7 @@ impl Entry { done, original_dtype: dtype, original_len: len, + reexecute_parent: false, } } } @@ -259,6 +268,7 @@ fn push_children( mut slots: Vec<(usize, DonePredicate)>, parent: Option, done: DonePredicate, + reexecute_parent: bool, stack: &mut Vec, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { @@ -289,7 +299,9 @@ fn push_children( ctx.log(format_args!( "ExecuteSlot({slot_idx}): pushing, focusing on {slot_child}" )); - stack.push(Entry::child(slot_child, parent_idx, slot_idx, slot_done)); + let mut entry = Entry::child(slot_child, parent_idx, slot_idx, slot_done); + entry.reexecute_parent = reexecute_parent; + stack.push(entry); } // Store parent back (with extracted slots emptied). From 8ac8c810daf8989371dec7a5992747d64c0f181a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 18:38:47 -0400 Subject: [PATCH 15/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 108 +++++++++++++++++++++++++++++++---- 1 file changed, 96 insertions(+), 12 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 22110d4012e..eadef277932 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -122,12 +122,45 @@ impl ArrayRef { current, &entry, )?; + if entry.run_execute_parent { + try_execute_parent_inline(&mut stack, link.parent_idx, ctx)?; + } continue; } } } - // Step 2: execute_parent, then execute encoding's own step. + // Step 2: if this child has run_execute_parent, try execute_parent + // on the reconstructed parent before executing the child further. + // This lets fused kernels fire (e.g. Filter(RunEnd) → RunEnd) + // before the child decodes to canonical. + let current = if entry.run_execute_parent { + if let Some(ref link) = entry.parent { + let parent = stack[link.parent_idx].array.take().unwrap(); + let parent = + unsafe { parent.put_slot_unchecked(link.slot_idx, current)? }; + match try_execute_parent_on(&parent, ctx)? { + Some(rewritten) => { + // Fused kernel fired — discard child, keep rewritten parent. + stack[link.parent_idx].array = Some(rewritten); + continue; + } + None => { + // No rewrite — take child back out, fall through to execute. + let (parent, child) = + unsafe { parent.take_slot_unchecked(link.slot_idx)? }; + stack[link.parent_idx].array = Some(parent); + child + } + } + } else { + current + } + } else { + current + }; + + // Step 3: execute_parent on the child's own children, then execute encoding's own step. let result = execute_one_step(current, ctx)?; match result { OneStepResult::Rewritten(rewritten) => { @@ -180,10 +213,13 @@ struct Entry { done: DonePredicate, original_dtype: DType, original_len: usize, - /// When true, after this child is put back into its parent, the parent is - /// popped and re-entered into the execute loop so that `execute_parent` can - /// fire on the reconstructed parent. Set for `ExecuteSlot` children only. - reexecute_parent: bool, + /// When true, after this child is put back into its parent, the scheduler + /// runs `execute_one_step` on the reconstructed parent so that + /// `execute_parent` can fire immediately. Set for children created by + /// [`ExecutionStep::ExecuteSlot`] (single child). Not set for children + /// created by [`ExecutionStep::ExecuteSlots`] (batch), where + /// `execute_parent` only fires after all children are done. + run_execute_parent: bool, } /// Points from a child [`Entry`] back to its parent in the stack. @@ -205,11 +241,17 @@ impl Entry { done, original_dtype: dtype, original_len: len, - reexecute_parent: false, + run_execute_parent: false, } } - fn child(array: ArrayRef, parent_idx: usize, slot_idx: usize, done: DonePredicate) -> Self { + fn child( + array: ArrayRef, + parent_idx: usize, + slot_idx: usize, + done: DonePredicate, + run_execute_parent: bool, + ) -> Self { let dtype = array.dtype().clone(); let len = array.len(); Self { @@ -221,7 +263,7 @@ impl Entry { done, original_dtype: dtype, original_len: len, - reexecute_parent: false, + run_execute_parent, } } } @@ -260,6 +302,50 @@ fn put_back_child( Ok(()) } +/// Try `execute_parent` on each child of the given array. Returns `Some` if any +/// child's `execute_parent` rewrites the parent. +fn try_execute_parent_on( + parent: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + for (slot_idx, slot) in parent.slots().iter().enumerate() { + let Some(child) = slot else { continue }; + if let Some(rewritten) = child.execute_parent(parent, slot_idx, ctx)? { + rewritten.statistics().inherit_from(parent.statistics()); + return Ok(Some(rewritten)); + } + } + Ok(None) +} + +/// After putting an ExecuteSlot child back, try `execute_parent` on each child +/// of the reconstructed parent. If any child's `execute_parent` rewrites the +/// parent, store the rewritten array. Otherwise leave the parent unchanged. +fn try_execute_parent_inline( + stack: &mut Vec, + parent_idx: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let parent = stack[parent_idx].array.take().unwrap(); + for (slot_idx, slot) in parent.slots().iter().enumerate() { + let Some(child) = slot else { continue }; + if let Some(rewritten) = child.execute_parent(&parent, slot_idx, ctx)? { + rewritten.statistics().inherit_from(parent.statistics()); + ctx.log(format_args!( + "execute_parent inline: slot[{}]({}) rewrote {} -> {}", + slot_idx, + child.encoding_id(), + parent, + rewritten + )); + stack[parent_idx].array = Some(rewritten); + return Ok(()); + } + } + stack[parent_idx].array = Some(parent); + Ok(()) +} + /// Extract all undone children from the parent and push them onto the stack in /// reverse order (first-to-process on top). Each child puts itself back into the /// parent via [`put_back_child`] when done. @@ -268,7 +354,7 @@ fn push_children( mut slots: Vec<(usize, DonePredicate)>, parent: Option, done: DonePredicate, - reexecute_parent: bool, + run_execute_parent: bool, stack: &mut Vec, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { @@ -299,9 +385,7 @@ fn push_children( ctx.log(format_args!( "ExecuteSlot({slot_idx}): pushing, focusing on {slot_child}" )); - let mut entry = Entry::child(slot_child, parent_idx, slot_idx, slot_done); - entry.reexecute_parent = reexecute_parent; - stack.push(entry); + stack.push(Entry::child(slot_child, parent_idx, slot_idx, slot_done, run_execute_parent)); } // Store parent back (with extracted slots emptied). From 9fac4611459884227afb5d6e69540bbae8804ad0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Apr 2026 18:44:21 -0400 Subject: [PATCH 16/16] fix Signed-off-by: Joe Isaacs --- vortex-array/src/executor.rs | 74 ++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index eadef277932..2bc80bc6952 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -130,34 +130,14 @@ impl ArrayRef { } } - // Step 2: if this child has run_execute_parent, try execute_parent - // on the reconstructed parent before executing the child further. - // This lets fused kernels fire (e.g. Filter(RunEnd) → RunEnd) - // before the child decodes to canonical. - let current = if entry.run_execute_parent { - if let Some(ref link) = entry.parent { - let parent = stack[link.parent_idx].array.take().unwrap(); - let parent = - unsafe { parent.put_slot_unchecked(link.slot_idx, current)? }; - match try_execute_parent_on(&parent, ctx)? { - Some(rewritten) => { - // Fused kernel fired — discard child, keep rewritten parent. - stack[link.parent_idx].array = Some(rewritten); - continue; - } - None => { - // No rewrite — take child back out, fall through to execute. - let (parent, child) = - unsafe { parent.take_slot_unchecked(link.slot_idx)? }; - stack[link.parent_idx].array = Some(parent); - child - } - } - } else { - current - } - } else { - current + // Step 2: for ExecuteSlot children, temporarily put the child back + // into its parent and try execute_parent. This lets fused kernels + // fire (e.g. Filter(RunEnd) → RunEnd) before the child decodes + // further. If a kernel fires, the parent is rewritten and the child + // is discarded. + let current = match try_fuse_with_parent(current, &entry, &mut stack, ctx)? { + FuseResult::Fused => continue, + FuseResult::Continue(child) => child, }; // Step 3: execute_parent on the child's own children, then execute encoding's own step. @@ -318,6 +298,44 @@ fn try_execute_parent_on( Ok(None) } +enum FuseResult { + /// A fused kernel fired — the parent was rewritten and the child discarded. + Fused, + /// No fusion — the child is returned for normal execution. + Continue(ArrayRef), +} + +/// For `ExecuteSlot` children (`run_execute_parent = true`), temporarily put the +/// child back into its parent and try `execute_parent`. If a fused kernel fires, +/// the parent entry is updated and the child is discarded. Otherwise the child is +/// taken back out and returned for normal execution. +fn try_fuse_with_parent( + child: ArrayRef, + entry: &Entry, + stack: &mut Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let Some(ref link) = entry.parent else { + return Ok(FuseResult::Continue(child)); + }; + if !entry.run_execute_parent { + return Ok(FuseResult::Continue(child)); + } + + let parent = stack[link.parent_idx].array.take().unwrap(); + let parent = unsafe { parent.put_slot_unchecked(link.slot_idx, child)? }; + + if let Some(rewritten) = try_execute_parent_on(&parent, ctx)? { + stack[link.parent_idx].array = Some(rewritten); + return Ok(FuseResult::Fused); + } + + // No rewrite — take child back out. + let (parent, child) = unsafe { parent.take_slot_unchecked(link.slot_idx)? }; + stack[link.parent_idx].array = Some(parent); + Ok(FuseResult::Continue(child)) +} + /// After putting an ExecuteSlot child back, try `execute_parent` on each child /// of the reconstructed parent. If any child's `execute_parent` rewrites the /// parent, store the rewritten array. Otherwise leave the parent unchanged.