From aa2789966acdb172c85b4f41de44787c29a323ee Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 15:19:17 -0400 Subject: [PATCH 01/19] Move execute-parent kernels into sessions Signed-off-by: "Nicholas Gates" Signed-off-by: Nicholas Gates --- encodings/alp/src/alp/array.rs | 10 -- encodings/alp/src/alp/mod.rs | 5 + encodings/alp/src/alp/rules.rs | 25 ++-- encodings/alp/src/alp_rd/array.rs | 10 -- encodings/alp/src/alp_rd/kernel.rs | 18 ++- encodings/alp/src/alp_rd/mod.rs | 5 + encodings/alp/src/lib.rs | 2 + encodings/bytebool/src/array.rs | 11 -- encodings/bytebool/src/kernel.rs | 16 ++- encodings/bytebool/src/lib.rs | 8 ++ encodings/datetime-parts/src/array.rs | 10 -- .../datetime-parts/src/compute/kernel.rs | 24 +++- encodings/datetime-parts/src/lib.rs | 1 + .../src/decimal_byte_parts/compute/kernel.rs | 24 +++- .../src/decimal_byte_parts/mod.rs | 10 -- encodings/decimal-byte-parts/src/lib.rs | 1 + encodings/experimental/onpair/src/array.rs | 10 -- encodings/experimental/onpair/src/kernel.rs | 23 +++- encodings/experimental/onpair/src/lib.rs | 8 ++ encodings/fastlanes/src/bitpacking/mod.rs | 4 + .../src/bitpacking/vtable/kernels.rs | 36 ++++-- .../fastlanes/src/bitpacking/vtable/mod.rs | 14 +-- .../src/bitpacking/vtable/operations.rs | 22 +++- encodings/fastlanes/src/for/mod.rs | 4 + encodings/fastlanes/src/for/vtable/kernels.rs | 16 ++- encodings/fastlanes/src/for/vtable/mod.rs | 14 +-- encodings/fastlanes/src/lib.rs | 3 + encodings/fastlanes/src/rle/kernel.rs | 12 +- encodings/fastlanes/src/rle/mod.rs | 4 + encodings/fastlanes/src/rle/vtable/mod.rs | 10 -- encodings/fsst/src/array.rs | 10 -- encodings/fsst/src/kernel.rs | 28 +++-- encodings/fsst/src/lib.rs | 8 ++ encodings/parquet-variant/src/kernel.rs | 33 +++-- encodings/parquet-variant/src/lib.rs | 1 + encodings/parquet-variant/src/vtable.rs | 10 -- encodings/runend/src/array.rs | 10 -- encodings/runend/src/kernel.rs | 23 ++-- encodings/runend/src/lib.rs | 1 + encodings/sequence/src/array.rs | 10 -- encodings/sequence/src/kernel.rs | 19 ++- encodings/sequence/src/lib.rs | 1 + encodings/sparse/src/kernel.rs | 28 +++-- encodings/sparse/src/lib.rs | 16 +-- encodings/zigzag/src/array.rs | 10 -- encodings/zigzag/src/kernel.rs | 12 +- encodings/zigzag/src/lib.rs | 8 ++ vortex-array/src/array/erased.rs | 9 -- vortex-array/src/array/mod.rs | 35 ------ vortex-array/src/array/vtable/mod.rs | 11 -- vortex-array/src/arrays/bool/mod.rs | 4 + vortex-array/src/arrays/bool/vtable/kernel.rs | 23 ++-- vortex-array/src/arrays/bool/vtable/mod.rs | 14 +-- .../src/arrays/chunked/compute/kernel.rs | 26 ++-- vortex-array/src/arrays/chunked/mod.rs | 4 + vortex-array/src/arrays/chunked/vtable/mod.rs | 10 -- vortex-array/src/arrays/decimal/mod.rs | 4 + .../src/arrays/decimal/vtable/kernel.rs | 23 ++-- vortex-array/src/arrays/decimal/vtable/mod.rs | 14 +-- vortex-array/src/arrays/dict/mod.rs | 4 + vortex-array/src/arrays/dict/vtable/kernel.rs | 19 ++- vortex-array/src/arrays/dict/vtable/mod.rs | 14 +-- vortex-array/src/arrays/extension/mod.rs | 4 + .../src/arrays/extension/vtable/kernel.rs | 21 +++- .../src/arrays/extension/vtable/mod.rs | 14 +-- .../src/arrays/filter/execute/take/tests.rs | 53 +++++--- vortex-array/src/arrays/filter/kernel.rs | 12 +- vortex-array/src/arrays/filter/mod.rs | 4 + vortex-array/src/arrays/filter/vtable.rs | 10 -- .../src/arrays/fixed_size_list/mod.rs | 4 + .../arrays/fixed_size_list/vtable/kernel.rs | 25 +++- .../src/arrays/fixed_size_list/vtable/mod.rs | 13 +- .../src/arrays/list/compute/kernels.rs | 20 ++- vortex-array/src/arrays/list/compute/mod.rs | 4 +- vortex-array/src/arrays/list/mod.rs | 4 + vortex-array/src/arrays/list/vtable/mod.rs | 10 -- vortex-array/src/arrays/listview/mod.rs | 4 + .../src/arrays/listview/vtable/kernel.rs | 16 ++- .../src/arrays/listview/vtable/mod.rs | 14 +-- vortex-array/src/arrays/mod.rs | 18 +++ vortex-array/src/arrays/patched/mod.rs | 4 + .../src/arrays/patched/vtable/kernels.rs | 17 ++- vortex-array/src/arrays/patched/vtable/mod.rs | 14 +-- vortex-array/src/arrays/primitive/mod.rs | 4 + .../src/arrays/primitive/vtable/kernel.rs | 34 ++++-- .../src/arrays/primitive/vtable/mod.rs | 14 +-- .../src/arrays/struct_/compute/cast.rs | 69 +++++------ vortex-array/src/arrays/struct_/mod.rs | 5 + .../src/arrays/struct_/vtable/kernel.rs | 13 +- vortex-array/src/arrays/struct_/vtable/mod.rs | 14 +-- vortex-array/src/arrays/varbin/mod.rs | 4 + .../src/arrays/varbin/vtable/kernel.rs | 23 ++-- vortex-array/src/arrays/varbin/vtable/mod.rs | 14 +-- vortex-array/src/arrays/varbinview/mod.rs | 4 + .../src/arrays/varbinview/vtable/kernel.rs | 20 ++- .../src/arrays/varbinview/vtable/mod.rs | 14 +-- vortex-array/src/arrays/variant/mod.rs | 5 + .../src/arrays/variant/vtable/kernel.rs | 11 +- vortex-array/src/arrays/variant/vtable/mod.rs | 14 +-- vortex-array/src/executor.rs | 29 +++-- vortex-array/src/kernel.rs | 109 +---------------- vortex-array/src/lib.rs | 11 +- vortex-array/src/optimizer/kernels.rs | 114 +++++++++++++++--- vortex-file/src/lib.rs | 20 +-- vortex/src/lib.rs | 1 + 105 files changed, 883 insertions(+), 750 deletions(-) diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index c85c0ea1b44..f0f87e2229d 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -48,7 +48,6 @@ use vortex_session::registry::CachedId; use crate::ALPFloat; use crate::alp::Exponents; use crate::alp::decompress::execute_decompress; -use crate::alp::rules::PARENT_KERNELS; use crate::alp::rules::RULES; /// A [`ALP`]-encoded Vortex array. @@ -188,15 +187,6 @@ impl VTable for ALP { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[array_slots(ALP)] diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 88375fe2ae6..a88afe8f55f 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -61,9 +61,14 @@ use vortex_array::dtype::NativePType; use vortex_array::scalar::PValue; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_session::VortexSession; const SAMPLE_SIZE: usize = 32; +pub(crate) fn initialize(session: &VortexSession) { + rules::initialize(session); +} + #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] pub struct Exponents { pub e: u8, diff --git a/encodings/alp/src/alp/rules.rs b/encodings/alp/src/alp/rules.rs index fe13682ccd8..3b888ae8269 100644 --- a/encodings/alp/src/alp/rules.rs +++ b/encodings/alp/src/alp/rules.rs @@ -1,26 +1,35 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::fns::mask::Mask; use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::VortexSession; 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) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), ALP, CompareExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Filter.id(), ALP, FilterExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Mask.id(), ALP, MaskExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Slice.id(), ALP, SliceExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Dict.id(), ALP, TakeExecuteAdaptor(ALP)); +} pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)), diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index eaf98851665..0a812fbc691 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -53,7 +53,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::alp_rd::kernel::PARENT_KERNELS; use crate::alp_rd::rules::RULES; use crate::alp_rd_decode; @@ -307,15 +306,6 @@ impl VTable for ALPRD { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The left (most significant) parts of the real-double encoded values. diff --git a/encodings/alp/src/alp_rd/kernel.rs b/encodings/alp/src/alp_rd/kernel.rs index 66721cbb1f9..73670163954 100644 --- a/encodings/alp/src/alp_rd/kernel.rs +++ b/encodings/alp/src/alp_rd/kernel.rs @@ -1,15 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_session::VortexSession; 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)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Slice.id(), ALPRD, SliceExecuteAdaptor(ALPRD)); + kernels.register_execute_parent_kernel(Filter.id(), ALPRD, FilterExecuteAdaptor(ALPRD)); + kernels.register_execute_parent_kernel(Dict.id(), ALPRD, TakeExecuteAdaptor(ALPRD)); +} diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index 5f2e703cc0a..a9ac7ca82fe 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -36,6 +36,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; +use vortex_session::VortexSession; use vortex_utils::aliases::hash_map::HashMap; use crate::match_each_alp_float_ptype; @@ -55,6 +56,10 @@ const CUT_LIMIT: usize = 16; const MAX_DICT_SIZE: u8 = 8; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + mod private { pub trait Sealed {} diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 7da676d6c6c..1286f3e6e2b 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -39,6 +39,8 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(ALP); } session.arrays().register(ALPRD); + alp::initialize(session); + alp_rd::initialize(session); // Register the ALP-specific NaN count aggregate kernel. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 225bad39ded..6b92bead37f 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -39,8 +39,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::kernel::PARENT_KERNELS; - /// A [`ByteBool`]-encoded Vortex array. pub type ByteBoolArray = Array; @@ -157,15 +155,6 @@ impl VTable for ByteBool { BoolArray::new(boolean_buffer, validity).into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The validity bitmap indicating which elements are non-null. diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index ddf9679b3d5..056c4b183c8 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -1,13 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_session::VortexSession; use crate::ByteBool; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(ByteBool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), ByteBool, CastExecuteAdaptor(ByteBool)); + kernels.register_execute_parent_kernel(Dict.id(), ByteBool, TakeExecuteAdaptor(ByteBool)); +} diff --git a/encodings/bytebool/src/lib.rs b/encodings/bytebool/src/lib.rs index 3258341cdb9..9d0d98e8c0a 100644 --- a/encodings/bytebool/src/lib.rs +++ b/encodings/bytebool/src/lib.rs @@ -41,9 +41,17 @@ //! [spec]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#bit-boolean pub use array::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; mod array; mod compute; mod kernel; mod rules; mod slice; + +/// Initialize bytebool encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(ByteBool); + kernel::initialize(session); +} diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index ddc5ee52981..64c6419a012 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -42,7 +42,6 @@ use vortex_session::registry::CachedId; use crate::TemporalParts; use crate::canonical::decode_to_temporal; -use crate::compute::kernel::PARENT_KERNELS; use crate::compute::rules::PARENT_RULES; use crate::split_temporal; @@ -200,15 +199,6 @@ impl VTable for DateTimeParts { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[array_slots(DateTimeParts)] diff --git a/encodings/datetime-parts/src/compute/kernel.rs b/encodings/datetime-parts/src/compute/kernel.rs index e07f1d00f38..4bd0cf2e1b3 100644 --- a/encodings/datetime-parts/src/compute/kernel.rs +++ b/encodings/datetime-parts/src/compute/kernel.rs @@ -1,13 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::DateTimeParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DateTimeParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DateTimeParts)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + DateTimeParts, + CompareExecuteAdaptor(DateTimeParts), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + DateTimeParts, + TakeExecuteAdaptor(DateTimeParts), + ); +} diff --git a/encodings/datetime-parts/src/lib.rs b/encodings/datetime-parts/src/lib.rs index e061d69bccd..cdc65a44a71 100644 --- a/encodings/datetime-parts/src/lib.rs +++ b/encodings/datetime-parts/src/lib.rs @@ -21,6 +21,7 @@ use vortex_session::VortexSession; /// Initialize datetime-parts encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(DateTimeParts); + compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( DateTimeParts.id(), 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..5e8d28e3526 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 @@ -1,13 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::DecimalByteParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DecimalByteParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DecimalByteParts)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + DecimalByteParts, + CompareExecuteAdaptor(DecimalByteParts), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + DecimalByteParts, + TakeExecuteAdaptor(DecimalByteParts), + ); +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 42135a9ecaa..1b5cba48ff3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -46,7 +46,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::decimal_byte_parts::compute::kernel::PARENT_KERNELS; use crate::decimal_byte_parts::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -165,15 +164,6 @@ impl VTable for DecimalByteParts { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { to_canonical_decimal(&array, ctx).map(ExecutionResult::done) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The most significant parts of the decimal values. diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 86566027afa..36a53c3a614 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -23,6 +23,7 @@ use vortex_session::VortexSession; /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(DecimalByteParts); + compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( DecimalByteParts.id(), diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index f9e18c36ac4..f4c8c1aed85 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -43,7 +43,6 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_onpair; use crate::canonical::onpair_decode_views; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// An [`OnPair`]-encoded Vortex array. @@ -499,15 +498,6 @@ impl VTable for OnPair { Ok(()) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index 7eb7b761b4f..8863d750a72 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -1,16 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Filter; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; +use vortex_session::VortexSession; use crate::OnPair; // TODO: implement ListExecute & TakeExecute for OnPair -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(OnPair)), - ParentKernelSet::lift(&CompareExecuteAdaptor(OnPair)), - ParentKernelSet::lift(&ByteLengthExecuteAdaptor(OnPair)), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Filter.id(), OnPair, FilterExecuteAdaptor(OnPair)); + kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); + kernels.register_execute_parent_kernel( + ByteLength.id(), + OnPair, + ByteLengthExecuteAdaptor(OnPair), + ); +} diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index 94c18b6dec8..11b22f63bc1 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -27,3 +27,11 @@ pub use onpair::Bits; pub use onpair::Config; pub use onpair::Error as OnPairError; pub use onpair::Threshold; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Initialize OnPair encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(OnPair); + kernel::initialize(session); +} diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index bc2304af657..efa0677a91e 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -19,3 +19,7 @@ mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index 87332f736a7..eb0dd9b7a23 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -1,21 +1,39 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_session::VortexSession; use crate::BitPacked; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&CastExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&CompareExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&FilterExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&SliceExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&TakeExecuteAdaptor(BitPacked)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + BitPacked, + BetweenExecuteAdaptor(BitPacked), + ); + kernels.register_execute_parent_kernel(Cast.id(), BitPacked, CastExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel( + Binary.id(), + BitPacked, + CompareExecuteAdaptor(BitPacked), + ); + kernels.register_execute_parent_kernel(Filter.id(), BitPacked, FilterExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel(Slice.id(), BitPacked, SliceExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel(Dict.id(), BitPacked, TakeExecuteAdaptor(BitPacked)); +} diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 52b536e02c5..9f8d41de014 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -48,7 +48,6 @@ use crate::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; use crate::bitpacking::array::PATCH_SLOTS; -use crate::bitpacking::vtable::kernels::PARENT_KERNELS; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; @@ -58,6 +57,10 @@ mod validity; /// A [`BitPacked`]-encoded Vortex array. pub type BitPackedArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + #[derive(Clone, prost::Message)] pub struct BitPackedMetadata { #[prost(uint32, tag = "1")] @@ -268,15 +271,6 @@ impl VTable for BitPacked { )) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 4c277163719..a41d3dcebca 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -31,6 +31,7 @@ impl OperationsVTable for BitPacked { #[cfg(test)] mod test { use std::ops::Range; + use std::sync::LazyLock; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -51,12 +52,19 @@ mod test { use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; + static FASTLANES_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() } @@ -154,12 +162,14 @@ mod test { // reads buffers. The slice range 0..64 excludes the patch at index 64, so the // resulting array should have no patches. let array_ref = array.into_array(); - let slice_array = SliceArray::new(array_ref.clone(), 0..64); - let sliced = array_ref - .execute_parent(&slice_array.into_array(), 0, &mut ctx) - .expect("execute_parent failed") - .expect("expected slice kernel to execute"); - let sliced_bp = sliced.as_::().into_owned(); + let slice_array = SliceArray::new(array_ref, 0..64); + let mut ctx = FASTLANES_SESSION.create_execution_ctx(); + let sliced_bp = slice_array + .into_array() + .execute::(&mut ctx) + .expect("slice execution failed") + .as_::() + .into_owned(); assert!(sliced_bp.patches().is_none()); } diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs index 8a70fe30499..a5d495e3578 100644 --- a/encodings/fastlanes/src/for/mod.rs +++ b/encodings/fastlanes/src/for/mod.rs @@ -10,3 +10,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::FoR; pub use vtable::FoRArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/encodings/fastlanes/src/for/vtable/kernels.rs b/encodings/fastlanes/src/for/vtable/kernels.rs index 331b13eceef..6c63292bb9f 100644 --- a/encodings/fastlanes/src/for/vtable/kernels.rs +++ b/encodings/fastlanes/src/for/vtable/kernels.rs @@ -1,13 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::FoR; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(FoR)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FoR)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), FoR, CompareExecuteAdaptor(FoR)); + kernels.register_execute_parent_kernel(Dict.id(), FoR, TakeExecuteAdaptor(FoR)); +} diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 037a858892e..d3a0bb84811 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -37,7 +37,6 @@ use crate::FoRData; use crate::r#for::array::FoRArrayExt; use crate::r#for::array::SLOT_NAMES; use crate::r#for::array::for_decompress::decompress; -use crate::r#for::vtable::kernels::PARENT_KERNELS; use crate::r#for::vtable::rules::PARENT_RULES; mod kernels; @@ -49,6 +48,10 @@ mod validity; /// A [`FoR`]-encoded Vortex array. pub type FoRArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + impl ArrayHash for FoRData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.reference.hash(state); @@ -150,15 +153,6 @@ impl VTable for FoR { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(decompress(&array, ctx)?.into_array())) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 85b6d0259a1..84d87649028 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -47,6 +47,9 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); + bitpacking::initialize(session); + r#for::initialize(session); + rle::initialize(session); // Register the encoding-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index f5b654706cd..6e89cd6045d 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -4,20 +4,26 @@ use std::ops::Range; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Slice; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; -pub(crate) static PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&SliceExecuteAdaptor(RLE))]); +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE)); +} impl SliceKernel for RLE { fn slice( diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 1c580683908..649d11f307c 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -11,3 +11,7 @@ mod kernel; mod vtable; pub use vtable::RLE; pub use vtable::RLEArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernel::initialize(session); +} diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index 82c236ae810..5cb50c82953 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -38,7 +38,6 @@ use crate::rle::array::SLOT_NAMES; use crate::rle::array::VALUES_IDX_OFFSETS_SLOT; use crate::rle::array::VALUES_SLOT; use crate::rle::array::rle_decompress::rle_decompress; -use crate::rle::kernel::PARENT_KERNELS; use crate::rle::vtable::rules::RULES; mod operations; @@ -193,15 +192,6 @@ impl VTable for RLE { Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( rle_decompress(&array, ctx)?.into_array(), diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 83360a69e64..9a737fb3057 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -57,7 +57,6 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_fsst; use crate::canonical::fsst_decode_views; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`FSST`]-encoded Vortex array. @@ -313,15 +312,6 @@ impl VTable for FSST { Ok(()) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index d45ee2978e3..942182097fa 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -1,24 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_session::VortexSession; use crate::FSST; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(FSST)), - ParentKernelSet::lift(&CompareExecuteAdaptor(FSST)), - ParentKernelSet::lift(&FilterExecuteAdaptor(FSST)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FSST)), - ParentKernelSet::lift(&LikeExecuteAdaptor(FSST)), - ParentKernelSet::lift(&ByteLengthExecuteAdaptor(FSST)), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), FSST, CastExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Binary.id(), FSST, CompareExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Filter.id(), FSST, FilterExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Dict.id(), FSST, TakeExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Like.id(), FSST, LikeExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(ByteLength.id(), FSST, ByteLengthExecuteAdaptor(FSST)); +} #[cfg(test)] mod tests { diff --git a/encodings/fsst/src/lib.rs b/encodings/fsst/src/lib.rs index 3305c0e66fc..70dcc705249 100644 --- a/encodings/fsst/src/lib.rs +++ b/encodings/fsst/src/lib.rs @@ -27,3 +27,11 @@ mod tests; pub use array::*; pub use compress::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Initialize FSST encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(FSST); + kernel::initialize(session); +} diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index a136b483678..424cb9d1d78 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -14,9 +14,13 @@ use parquet_variant_compute::GetOptions; use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::variant_get as arrow_variant_get; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; @@ -28,7 +32,8 @@ use vortex_array::arrays::slice::SliceKernel; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::kernel::ExecuteParentKernel; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::variant_get::VariantGet; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; @@ -36,16 +41,30 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_mask::Mask; +use vortex_session::VortexSession; 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)), - ParentKernelSet::lift(&VariantGetKernel), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Filter.id(), + ParquetVariant, + FilterExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel( + Slice.id(), + ParquetVariant, + SliceExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + ParquetVariant, + TakeExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); +} #[derive(Default, Debug)] struct VariantGetKernel; diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 03d2a046442..0076af48d22 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -43,6 +43,7 @@ pub use vtable::ParquetVariantArray; /// Register Parquet Variant array and Arrow extension support with a session. pub fn initialize(session: &VortexSession) { session.arrays().register(ParquetVariant); + kernel::initialize(session); session.arrow().register_exporter(Arc::new(ParquetVariant)); session.arrow().register_importer(Arc::new(ParquetVariant)); } diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index d73071fc3a6..d5b98ce03aa 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -39,7 +39,6 @@ use crate::array::VALIDITY_SLOT; use crate::array::VALUE_SLOT; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; -use crate::kernel::PARENT_KERNELS; /// VTable for Arrow's canonical `arrow.parquet.variant` extension storage. /// @@ -280,15 +279,6 @@ impl VTable for ParquetVariant { VariantArray::try_new(core_storage, shredded)?.into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[cfg(test)] diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 99cecddc0b0..f69e9175222 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -48,7 +48,6 @@ use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -172,15 +171,6 @@ impl VTable for RunEnd { RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { run_end_canonicalize(&array, ctx).map(ExecutionResult::done) } diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 73f6ebc524e..484a4233789 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -4,29 +4,36 @@ use std::ops::Range; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; 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::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_error::VortexResult; +use vortex_session::VortexSession; 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), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), RunEnd, CompareExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Slice.id(), RunEnd, RunEndSliceKernel); + kernels.register_execute_parent_kernel(Filter.id(), RunEnd, FilterExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Dict.id(), RunEnd, TakeExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Dict.id(), RunEnd, RunEndTakeFrom); +} /// Kernel to execute slicing on a RunEnd array. /// diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8770dbbf58e..846c46d05d4 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -38,6 +38,7 @@ use vortex_session::VortexSession; /// Initialize run-end encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(RunEnd); + kernel::initialize(session); // Register the RunEnd-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index af8b146cfba..15365b04464 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -49,7 +49,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::compress::sequence_decompress; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`Sequence`]-encoded Vortex array. @@ -331,15 +330,6 @@ impl VTable for Sequence { sequence_decompress(&array).map(ExecutionResult::done) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/sequence/src/kernel.rs b/encodings/sequence/src/kernel.rs index 394a52ec69e..37521a261e7 100644 --- a/encodings/sequence/src/kernel.rs +++ b/encodings/sequence/src/kernel.rs @@ -1,15 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::Sequence; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&FilterExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sequence)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Sequence, CompareExecuteAdaptor(Sequence)); + kernels.register_execute_parent_kernel(Filter.id(), Sequence, FilterExecuteAdaptor(Sequence)); + kernels.register_execute_parent_kernel(Dict.id(), Sequence, TakeExecuteAdaptor(Sequence)); +} diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index bd6ab2f508c..b898963c346 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -27,6 +27,7 @@ use vortex_session::VortexSession; /// Initialize sequence encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(Sequence); + kernel::initialize(session); // Register the Sequence-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/sparse/src/kernel.rs b/encodings/sparse/src/kernel.rs index 0f5d9fd51c0..30e1f9a6d6a 100644 --- a/encodings/sparse/src/kernel.rs +++ b/encodings/sparse/src/kernel.rs @@ -1,21 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::fill_null::FillNull; use vortex_array::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use vortex_session::VortexSession; use crate::Sparse; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&CompareExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&FilterExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&SliceExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sparse)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Between.id(), Sparse, BetweenExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Binary.id(), Sparse, CompareExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(FillNull.id(), Sparse, FillNullExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Filter.id(), Sparse, FilterExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Slice.id(), Sparse, SliceExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Dict.id(), Sparse, TakeExecuteAdaptor(Sparse)); +} diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index e74cedb942e..d47cf36563e 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message as _; use vortex_array::AnyCanonical; use vortex_array::Array; @@ -79,11 +78,11 @@ use vortex_array::session::ArraySessionExt; /// Initialize Sparse encoding in the given session. /// -/// Registers the Sparse array vtable and its aggregate kernels (`IsConstant`, `Sum`, -/// `MinMax`, `NullCount`, `NanCount`). Compare/between/fill_null pushdown is wired -/// through `PARENT_KERNELS` (see `kernel.rs`) and does not require registration here. +/// Registers the Sparse array vtable, parent execution kernels, and aggregate kernels +/// (`IsConstant`, `Sum`, `MinMax`, `NullCount`, `NanCount`). pub fn initialize(session: &VortexSession) { session.arrays().register(Sparse); + kernel::initialize(session); let aggregate_fns = session.aggregate_fns(); aggregate_fns.register_aggregate_kernel( @@ -294,15 +293,6 @@ impl VTable for Sparse { RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { // Resolve offset first: wrap indices in Binary(indices, offset, Sub) and // reassemble with offset=0. Uses slot children (not data) since the executor diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index f03315ae575..281c23ca4c9 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -38,7 +38,6 @@ use vortex_session::registry::CachedId; use zigzag::ZigZag as ExternalZigZag; use crate::compute::ZigZagEncoded; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; use crate::zigzag_decode; @@ -143,15 +142,6 @@ impl VTable for ZigZag { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } impl ArrayHash for ZigZagData { diff --git a/encodings/zigzag/src/kernel.rs b/encodings/zigzag/src/kernel.rs index d0096abaae1..f1d3e96bb01 100644 --- a/encodings/zigzag/src/kernel.rs +++ b/encodings/zigzag/src/kernel.rs @@ -1,10 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_session::VortexSession; use crate::ZigZag; -pub(crate) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(ZigZag))]); +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Dict.id(), ZigZag, TakeExecuteAdaptor(ZigZag)); +} diff --git a/encodings/zigzag/src/lib.rs b/encodings/zigzag/src/lib.rs index 89da8bd6069..01c3f51541b 100644 --- a/encodings/zigzag/src/lib.rs +++ b/encodings/zigzag/src/lib.rs @@ -3,6 +3,8 @@ pub use array::*; pub use compress::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; mod array; mod compress; @@ -10,3 +12,9 @@ mod compute; mod kernel; mod rules; mod slice; + +/// Initialize zigzag encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(ZigZag); + kernel::initialize(session); +} diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index bacda957660..9b6d92014c9 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -623,15 +623,6 @@ impl ArrayRef { unsafe { (&*inner).data.execute_unchecked(self, ctx) } } - pub fn execute_parent( - &self, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - self.0.data.execute_parent(self, parent, child_idx, ctx) - } - // ArrayVisitor delegation methods /// Returns the children of the array. diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index b1b22c4c529..6e149d3a81a 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -177,15 +177,6 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult; - /// Attempt to execute the parent of this array. - fn execute_parent( - &self, - this: &ArrayRef, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; - /// Execute the scalar at the given index. /// /// This method panics if the index is out of bounds for the array. @@ -474,32 +465,6 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } - fn execute_parent( - &self, - this: &ArrayRef, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - let Some(result) = V::execute_parent(view, parent, child_idx, ctx)? else { - return Ok(None); - }; - - if cfg!(debug_assertions) { - vortex_ensure!( - result.len() == parent.len(), - "Executed parent canonical length mismatch" - ); - vortex_ensure!( - result.dtype() == parent.dtype(), - "Executed parent canonical dtype mismatch" - ); - } - - Ok(Some(result)) - } - fn execute_scalar( &self, this: &ArrayRef, diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 366eed77ad9..890290116e7 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -185,17 +185,6 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { /// incorrectly contains null values. fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult; - /// Attempt to execute the parent of this array. - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - _ = (array, parent, child_idx, ctx); - Ok(None) - } - /// Attempt to reduce the array to a simpler representation. fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { _ = array; diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index a4aee2c4faa..ea51e6b5f29 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -15,5 +15,9 @@ pub use compute::rules::BoolMaskedValidityRule; pub use vtable::Bool; pub use vtable::BoolArray; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(feature = "_test-harness")] mod test_harness; diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index 2fe2d974a8f..fb627a34846 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -1,16 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Bool; +use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(Bool)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Bool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Bool)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Bool)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), Bool, CastExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(FillNull.id(), Bool, FillNullExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Dict.id(), Bool, TakeExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Zip.id(), Bool, ZipExecuteAdaptor(Bool)); +} diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 1bbc091b010..264c3ab3a3b 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -43,6 +42,10 @@ use crate::hash::ArrayHash; /// A [`Bool`]-encoded Vortex array. pub type BoolArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(prost::Message)] pub struct BoolMetadata { // The offset in bits must be <8 @@ -174,15 +177,6 @@ impl VTable for Bool { Ok(ExecutionResult::done(array)) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/chunked/compute/kernel.rs b/vortex-array/src/arrays/chunked/compute/kernel.rs index 502c1fe8337..db0042105cd 100644 --- a/vortex-array/src/arrays/chunked/compute/kernel.rs +++ b/vortex-array/src/arrays/chunked/compute/kernel.rs @@ -1,18 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Chunked; +use crate::arrays::Dict; +use crate::arrays::Filter; +use crate::arrays::Slice; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; 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)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Filter.id(), Chunked, FilterExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Mask.id(), Chunked, MaskExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Slice.id(), Chunked, SliceExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Dict.id(), Chunked, TakeExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Zip.id(), Chunked, ZipExecuteAdaptor(Chunked)); +} diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 766d9f6575c..4bd6cb43b7e 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -12,5 +12,9 @@ pub(crate) mod paired_chunks; mod vtable; pub use vtable::Chunked; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::kernel::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index a70ba6869ab..1054dd21d2d 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -33,7 +33,6 @@ use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::CHUNK_OFFSETS_SLOT; use crate::arrays::chunked::array::CHUNKS_OFFSET; -use crate::arrays::chunked::compute::kernel::PARENT_KERNELS; use crate::arrays::chunked::compute::rules::PARENT_RULES; use crate::arrays::chunked::vtable::canonical::_canonicalize; use crate::buffer::BufferHandle; @@ -263,15 +262,6 @@ impl VTable for Chunked { } } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { Ok(match array.nchunks() { 0 => Some(Canonical::empty(array.dtype()).into_array()), diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index 489a67a0531..e740a75ea91 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -13,6 +13,10 @@ mod vtable; pub use compute::rules::DecimalMaskedValidityRule; pub use vtable::Decimal; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + mod utils; pub use utils::*; diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 9ad3803be1d..19bb30d8c95 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -1,16 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Decimal; +use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; 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)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Between.id(), Decimal, BetweenExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(Cast.id(), Decimal, CastExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(FillNull.id(), Decimal, FillNullExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(Dict.id(), Decimal, TakeExecuteAdaptor(Decimal)); +} diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 74e76ab4f51..a8dce4a1254 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_buffer::Alignment; use vortex_error::VortexResult; @@ -43,6 +42,10 @@ use crate::hash::ArrayHash; /// A [`Decimal`]-encoded Vortex array. pub type DecimalArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + // The type of the values can be determined by looking at the type info...right? #[derive(prost::Message)] pub struct DecimalMetadata { @@ -192,15 +195,6 @@ impl VTable for Decimal { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/dict/mod.rs b/vortex-array/src/arrays/dict/mod.rs index 0414eea7def..e8a83aa88b3 100644 --- a/vortex-array/src/arrays/dict/mod.rs +++ b/vortex-array/src/arrays/dict/mod.rs @@ -23,5 +23,9 @@ pub use take::*; pub mod vtable; pub use vtable::*; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/dict/vtable/kernel.rs b/vortex-array/src/arrays/dict/vtable/kernel.rs index fd1b0ed0a6d..ab750f7d663 100644 --- a/vortex-array/src/arrays/dict/vtable/kernel.rs +++ b/vortex-array/src/arrays/dict/vtable/kernel.rs @@ -1,14 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; 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)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Dict, CompareExecuteAdaptor(Dict)); + kernels.register_execute_parent_kernel(Dict.id(), Dict, TakeExecuteAdaptor(Dict)); + kernels.register_execute_parent_kernel(FillNull.id(), Dict, FillNullExecuteAdaptor(Dict)); +} diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index eb742182a10..bb9eb195132 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use smallvec::smallvec; use vortex_error::VortexResult; @@ -54,6 +53,10 @@ mod validity; /// A [`Dict`]-encoded Vortex array. pub type DictArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct Dict; @@ -203,13 +206,4 @@ impl VTable for Dict { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } diff --git a/vortex-array/src/arrays/extension/mod.rs b/vortex-array/src/arrays/extension/mod.rs index 57bc9df5ead..8dc5972b86f 100644 --- a/vortex-array/src/arrays/extension/mod.rs +++ b/vortex-array/src/arrays/extension/mod.rs @@ -9,3 +9,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::Extension; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index 4cc5f9eb4ac..c90458f1d5f 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -1,12 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; use crate::arrays::Extension; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Extension)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Extension)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + Extension, + CompareExecuteAdaptor(Extension), + ); + kernels.register_execute_parent_kernel(Dict.id(), Extension, TakeExecuteAdaptor(Extension)); +} diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index 852593abddb..a1fab1cd47a 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use kernel::PARENT_KERNELS; use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -78,6 +77,10 @@ pub struct Extension; /// A [`Extension`]-encoded Vortex array. pub type ExtensionArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl VTable for Extension { type TypedArrayData = EmptyArrayData; @@ -176,15 +179,6 @@ impl VTable for Extension { Ok(ExecutionResult::done(array)) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { RULES.evaluate(array) } diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index 3d3a0d1a4e3..96728f988a6 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -21,8 +21,30 @@ use crate::assert_arrays_eq; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::executor::ExecutionCtx; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::validity::Validity; +fn execute_parent( + child: &crate::ArrayRef, + parent: &crate::ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let kernels = ctx.session().kernels().clone(); + let Some(plugins) = kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) + else { + return Ok(None); + }; + + for plugin in plugins.as_ref() { + if let Some(result) = plugin.execute_parent(child, parent, child_idx, ctx)? { + return Ok(Some(result)); + } + } + + Ok(None) +} + #[test] fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { let filter = FilterArray::new( @@ -42,8 +64,7 @@ fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { .into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!( @@ -71,8 +92,7 @@ fn test_take_execute_kernel_nullable_fast_path_maps_indices_through_filter() -> .into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert!(result.as_opt::().is_some()); @@ -93,8 +113,7 @@ fn test_take_execute_kernel_fast_path_maps_indices_through_filter() -> VortexRes let parent = DictArray::try_new(buffer![2u64, 0, 3].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert!(result.as_opt::().is_some()); @@ -114,7 +133,7 @@ fn assert_take_execute_rejects_out_of_bounds_rank( let parent = DictArray::try_new(codes, filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - if let Err(err) = filter.execute_parent(&parent, 1, &mut ctx) { + if let Err(err) = execute_parent(&filter, &parent, 1, &mut ctx) { assert!( err.to_string().contains("out of bounds"), "unexpected error: {err}" @@ -192,8 +211,7 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> .into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!( @@ -217,8 +235,7 @@ fn assert_take_execute_maps_child_dtype( let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!(result.execute::(&mut ctx)?.0, expected); @@ -235,7 +252,7 @@ fn test_take_execute_kernel_skips_bool_filter_child() -> VortexResult<()> { let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert!(result.is_none()); Ok(()) @@ -257,7 +274,7 @@ fn execute_primitive_take( let parent = DictArray::try_new(indices.into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - filter.execute_parent(&parent, 1, &mut ctx) + execute_parent(&filter, &parent, 1, &mut ctx) } #[test] @@ -324,7 +341,7 @@ fn test_take_execute_kernel_handles_nullable_primitive_filter_child() -> VortexR let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert_arrays_eq!( result @@ -346,8 +363,7 @@ fn test_take_execute_kernel_preserves_nullable_all_valid_fixed_width_child() -> let parent = DictArray::try_new(buffer![0u64, 1].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_eq!(result.dtype(), parent.dtype()); @@ -373,7 +389,7 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert_arrays_eq!( result @@ -466,8 +482,7 @@ fn test_take_execute_kernel_preserves_nullable_indices_dtype_fast_path() -> Vort .into_array(); let mut ctx = ExecutionCtx::new(crate::array_session()); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its nullable take parent"); assert_eq!(result.dtype(), parent.dtype()); diff --git a/vortex-array/src/arrays/filter/kernel.rs b/vortex-array/src/arrays/filter/kernel.rs index 8c65fa99724..21bd225bf55 100644 --- a/vortex-array/src/arrays/filter/kernel.rs +++ b/vortex-array/src/arrays/filter/kernel.rs @@ -10,22 +10,28 @@ use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::VortexSession; use crate::ArrayRef; +use crate::ArrayVTable; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Dict; use crate::arrays::Filter; use crate::arrays::dict::TakeExecuteAdaptor; use crate::kernel::ExecuteParentKernel; -use crate::kernel::ParentKernelSet; use crate::matcher::Matcher; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::rules::ArrayParentReduceRule; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(Filter))]); +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter)); +} pub trait FilterReduce: VTable { /// Filter an array with the provided mask without reading buffers. diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index 39859bc5b25..9510ca92f5b 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -19,3 +19,7 @@ mod rules; mod vtable; pub use vtable::Filter; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernel::initialize(session); +} diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 5be6bc6ac35..af11393c5db 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -32,7 +32,6 @@ use crate::arrays::filter::array::FilterData; use crate::arrays::filter::array::SLOT_NAMES; use crate::arrays::filter::execute::execute_filter; use crate::arrays::filter::execute::execute_filter_fast_paths; -use crate::arrays::filter::kernel::PARENT_KERNELS; use crate::arrays::filter::rules::PARENT_RULES; use crate::arrays::filter::rules::RULES; use crate::buffer::BufferHandle; @@ -171,15 +170,6 @@ impl VTable for Filter { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { RULES.evaluate(array) } diff --git a/vortex-array/src/arrays/fixed_size_list/mod.rs b/vortex-array/src/arrays/fixed_size_list/mod.rs index 5d55d22d880..3d90d38e6e9 100644 --- a/vortex-array/src/arrays/fixed_size_list/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/mod.rs @@ -12,5 +12,9 @@ pub(crate) mod compute; mod vtable; pub use vtable::FixedSizeList; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(test)] mod tests; 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 319751070de..47dc3ec0651 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -1,14 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; use crate::arrays::FixedSizeList; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -impl FixedSizeList { - pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(FixedSizeList)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FixedSizeList)), - ]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Cast.id(), + FixedSizeList, + CastExecuteAdaptor(FixedSizeList), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + FixedSizeList, + TakeExecuteAdaptor(FixedSizeList), + ); } 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 8d86e96a174..bfc172a5049 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -39,6 +39,10 @@ mod validity; /// A [`FixedSizeList`]-encoded Vortex array. pub type FixedSizeListArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct FixedSizeList; @@ -84,15 +88,6 @@ impl VTable for FixedSizeList { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Self::PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index 188c83c1bf5..6af858dbf49 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -1,14 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; +use crate::arrays::Filter; use crate::arrays::List; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(List)), - ParentKernelSet::lift(&FilterExecuteAdaptor(List)), - ParentKernelSet::lift(&TakeExecuteAdaptor(List)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), List, CastExecuteAdaptor(List)); + kernels.register_execute_parent_kernel(Filter.id(), List, FilterExecuteAdaptor(List)); + kernels.register_execute_parent_kernel(Dict.id(), List, TakeExecuteAdaptor(List)); +} diff --git a/vortex-array/src/arrays/list/compute/mod.rs b/vortex-array/src/arrays/list/compute/mod.rs index 65bacbebdf5..ac81556d050 100644 --- a/vortex-array/src/arrays/list/compute/mod.rs +++ b/vortex-array/src/arrays/list/compute/mod.rs @@ -9,7 +9,9 @@ pub(crate) mod rules; mod slice; mod take; -pub(crate) use kernels::PARENT_KERNELS; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernels::initialize(session); +} #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index c884f343a10..422a9523ff9 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -12,6 +12,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::List; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::initialize(session); +} + #[cfg(feature = "_test-harness")] mod test_harness; diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index d93c4673302..c71adadd8a0 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -31,7 +31,6 @@ use crate::arrays::list::array::ELEMENTS_SLOT; use crate::arrays::list::array::NUM_SLOTS; use crate::arrays::list::array::OFFSETS_SLOT; use crate::arrays::list::array::SLOT_NAMES; -use crate::arrays::list::compute::PARENT_KERNELS; use crate::arrays::list::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; @@ -191,15 +190,6 @@ impl VTable for List { list_view_from_list(array, ctx)?.into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index d5674372b2c..80b2540a897 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -12,6 +12,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::ListView; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + mod conversion; pub use conversion::list_from_list_view; pub use conversion::list_view_from_list; diff --git a/vortex-array/src/arrays/listview/vtable/kernel.rs b/vortex-array/src/arrays/listview/vtable/kernel.rs index 1ad98f62a33..5621b23e79e 100644 --- a/vortex-array/src/arrays/listview/vtable/kernel.rs +++ b/vortex-array/src/arrays/listview/vtable/kernel.rs @@ -1,12 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + use crate::arrays::ListView; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(ListView)), - ParentKernelSet::lift(&ZipExecuteAdaptor(ListView)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), ListView, CastExecuteAdaptor(ListView)); + kernels.register_execute_parent_kernel(Zip.id(), ListView, ZipExecuteAdaptor(ListView)); +} diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 26e54b2d051..ba5ae42b3ad 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -32,7 +32,6 @@ use crate::arrays::listview::array::OFFSETS_SLOT; use crate::arrays::listview::array::SIZES_SLOT; use crate::arrays::listview::array::SLOT_NAMES; use crate::arrays::listview::compute::rules::PARENT_RULES; -use crate::arrays::listview::vtable::kernel::PARENT_KERNELS; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; @@ -45,6 +44,10 @@ mod validity; /// A [`ListView`]-encoded Vortex array. pub type ListViewArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct ListView; @@ -219,13 +222,4 @@ impl VTable for ListView { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 148cf209555..64a10c4785b 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -106,5 +106,23 @@ pub mod variant; pub use variant::Variant; pub use variant::VariantArray; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + bool::initialize(session); + chunked::initialize(session); + decimal::initialize(session); + dict::initialize(session); + extension::initialize(session); + filter::initialize(session); + fixed_size_list::initialize(session); + list::initialize(session); + listview::initialize(session); + patched::initialize(session); + primitive::initialize(session); + struct_::initialize(session); + varbin::initialize(session); + varbinview::initialize(session); + variant::initialize(session); +} + #[cfg(feature = "arbitrary")] pub mod arbitrary; diff --git a/vortex-array/src/arrays/patched/mod.rs b/vortex-array/src/arrays/patched/mod.rs index 7d6acd395e7..a88628e5587 100644 --- a/vortex-array/src/arrays/patched/mod.rs +++ b/vortex-array/src/arrays/patched/mod.rs @@ -78,6 +78,10 @@ pub use array::*; use vortex_buffer::ByteBuffer; pub use vtable::*; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + /// Patches that have been transposed into GPU format. struct TransposedPatches { n_lanes: usize, diff --git a/vortex-array/src/arrays/patched/vtable/kernels.rs b/vortex-array/src/arrays/patched/vtable/kernels.rs index 7994b19e02e..51f53d15625 100644 --- a/vortex-array/src/arrays/patched/vtable/kernels.rs +++ b/vortex-array/src/arrays/patched/vtable/kernels.rs @@ -1,12 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; use crate::arrays::Patched; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Patched)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Patched)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Patched, CompareExecuteAdaptor(Patched)); + kernels.register_execute_parent_kernel(Dict.id(), Patched, TakeExecuteAdaptor(Patched)); +} diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 4ba5bc542e1..7790d734ed6 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -40,7 +40,6 @@ use crate::arrays::patched::PatchedData; use crate::arrays::patched::PatchedSlots; use crate::arrays::patched::PatchedSlotsView; use crate::arrays::patched::compute::rules::PARENT_RULES; -use crate::arrays::patched::vtable::kernels::PARENT_KERNELS; use crate::arrays::primitive::PrimitiveDataParts; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -55,6 +54,10 @@ use crate::serde::ArrayChildren; /// A [`Patched`]-encoded Vortex array. pub type PatchedArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + #[derive(Clone, Debug)] pub struct Patched; @@ -301,15 +304,6 @@ impl VTable for Patched { Ok(ExecutionResult::done(patched_values.into_array())) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index 4d62f1da517..e291ff521fb 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -15,6 +15,10 @@ mod vtable; pub use compute::rules::PrimitiveMaskedValidityRule; pub use vtable::Primitive; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + mod native_value; pub use native_value::NativeValue; diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 341d2f21272..6382ea73794 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -1,18 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; use crate::arrays::Primitive; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&CastExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Primitive)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + Primitive, + BetweenExecuteAdaptor(Primitive), + ); + kernels.register_execute_parent_kernel(Cast.id(), Primitive, CastExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel( + FillNull.id(), + Primitive, + FillNullExecuteAdaptor(Primitive), + ); + kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); +} diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 4b472d122f3..90f56ea1aa5 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use kernel::PARENT_KERNELS; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -39,6 +38,10 @@ use crate::hash::ArrayHash; /// A [`Primitive`]-encoded Vortex array. pub type PrimitiveArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl ArrayHash for PrimitiveData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.buffer.array_hash(state, accuracy); @@ -191,15 +194,6 @@ impl VTable for Primitive { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index 1f8ee705e9d..4ca1febe851 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::ArrayView; @@ -13,33 +14,43 @@ use crate::arrays::ConstantArray; use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::StructFields; -use crate::matcher::Matcher; +use crate::kernel::ExecuteParentKernel; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; -pub(crate) fn struct_cast_execute_parent( - child: &ArrayRef, - parent: &ArrayRef, - _child_idx: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some(array) = child.as_opt::() else { - return Ok(None); - }; - let Some(parent) = ExactScalarFn::::try_match(parent) else { - return Ok(None); - }; +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Cast.id(), Struct, StructCastKernel); +} - let dtype = parent.options; - if array.dtype() == parent.options { - return Ok(Some(array.array().clone())); - } +#[derive(Debug)] +struct StructCastKernel; + +impl ExecuteParentKernel for StructCastKernel { + type Parent = ExactScalarFn; - struct_cast(array, dtype, ctx) + fn execute_parent( + &self, + array: ArrayView<'_, Struct>, + parent: ScalarFnArrayView<'_, Cast>, + _child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let dtype = parent.options; + if array.dtype() == parent.options { + return Ok(Some(array.array().clone())); + } + + struct_cast(array, dtype, ctx) + } } pub(crate) fn struct_cast( @@ -140,6 +151,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::kernels::ExecuteParentFn; use crate::scalar::Scalar; @@ -180,25 +192,6 @@ mod tests { test_cast_conformance(&array.into_array()); } - #[test] - fn struct_cast_execute_parent_is_not_static_kernel() { - let source = create_simple_struct().into_array(); - let target = DType::struct_( - [( - "value", - DType::Primitive(PType::I64, Nullability::NonNullable), - )], - Nullability::NonNullable, - ); - - let cast = Cast - .try_new_array(source.len(), target, [source.clone()]) - .unwrap(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); - - assert!(source.execute_parent(&cast, 0, &mut ctx).unwrap().is_none()); - } - #[test] fn struct_cast_execute_parent_uses_session_plugin() { let source = StructArray::try_new( @@ -221,7 +214,7 @@ mod tests { .try_new_array(source.len(), target.clone(), [source]) .unwrap(); let parent_id = cast.encoding_id(); - let session = crate::array_session(); + let session = VortexSession::empty().with::(); session.kernels().register_execute_parent( parent_id, child_id, diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index 5fccb58f9cd..da49411c3bc 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -10,5 +10,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::Struct; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::cast::initialize(session); + vtable::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index eac7158921d..aa7f648b935 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -1,9 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + use crate::arrays::Struct; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&ZipExecuteAdaptor(Struct))]); +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); +} diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 71ec8a8e2df..df9b1170b02 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use kernel::PARENT_KERNELS; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -37,6 +36,10 @@ use crate::array::ArrayId; /// A [`Struct`]-encoded Vortex array. pub type StructArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl VTable for Struct { type TypedArrayData = EmptyArrayData; @@ -192,15 +195,6 @@ impl VTable for Struct { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/varbin/mod.rs b/vortex-array/src/arrays/varbin/mod.rs index a8f321fec9a..f0ccbe4585c 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/vortex-array/src/arrays/varbin/mod.rs @@ -12,6 +12,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBin; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + pub mod builder; mod accessor; diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 9258e677933..9e80abd1037 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -1,16 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; +use crate::arrays::Filter; use crate::arrays::VarBin; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&CompareExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&FilterExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBin)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), VarBin, CastExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Binary.id(), VarBin, CompareExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Filter.id(), VarBin, FilterExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Dict.id(), VarBin, TakeExecuteAdaptor(VarBin)); +} diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index dae80b663de..71003aa9cad 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -36,7 +36,6 @@ mod operations; mod validity; use canonical::varbin_to_canonical; -use kernel::PARENT_KERNELS; use vortex_session::VortexSession; use crate::EqMode; @@ -47,6 +46,10 @@ use crate::hash::ArrayHash; /// A [`VarBin`]-encoded Vortex array. pub type VarBinArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, prost::Message)] pub struct VarBinMetadata { #[prost(enumeration = "PType", tag = "1")] @@ -182,15 +185,6 @@ impl VTable for VarBin { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( varbin_to_canonical(array.as_view(), ctx)?.into_array(), diff --git a/vortex-array/src/arrays/varbinview/mod.rs b/vortex-array/src/arrays/varbinview/mod.rs index be82b952f30..6f3398884b3 100644 --- a/vortex-array/src/arrays/varbinview/mod.rs +++ b/vortex-array/src/arrays/varbinview/mod.rs @@ -15,6 +15,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBinView; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + pub mod build_views; mod view; diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index cd9d68010af..e09b381d590 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -1,14 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(VarBinView)), - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBinView)), - ParentKernelSet::lift(&ZipExecuteAdaptor(VarBinView)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); +} diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index c093b9c0203..d6d6b98aa7e 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hasher; use std::mem::size_of; use std::sync::Arc; -use kernel::PARENT_KERNELS; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -40,6 +39,10 @@ mod validity; /// A [`VarBinView`]-encoded Vortex array. pub type VarBinViewArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct VarBinView; @@ -217,15 +220,6 @@ impl VTable for VarBinView { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 7d7f23e34c5..a04894f2981 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -11,6 +11,11 @@ use vortex_error::vortex_ensure; pub use self::vtable::Variant; pub use self::vtable::VariantArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-array/src/arrays/variant/vtable/kernel.rs index a24eb81fd5d..bfe89dea810 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-array/src/arrays/variant/vtable/kernel.rs @@ -3,6 +3,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::VortexSession; use super::merge_typed_scalar_as_variant; use crate::ArrayRef; @@ -24,14 +25,18 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::variant_get::VariantGet; use crate::scalar_fn::fns::variant_get::VariantGetOptions; use crate::scalar_fn::fns::variant_get::VariantPath; use crate::scalar_fn::fns::variant_get::VariantPathElement; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&VariantGetKernel)]); +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(VariantGet.id(), Variant, VariantGetKernel); +} #[derive(Default, Debug)] struct VariantGetKernel; diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 3c0dd76c575..68c12cb634e 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -5,7 +5,6 @@ mod kernel; mod operations; mod validity; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -43,6 +42,10 @@ use crate::serde::ArrayChildren; /// A [`Variant`]-encoded Vortex array. pub type VariantArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct Variant; @@ -192,15 +195,6 @@ impl VTable for Variant { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } fn merge_typed_scalar_as_variant( diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 2bd5dfd5b37..c2f5eda9424 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -138,12 +138,16 @@ impl ArrayRef { /// - yes -> skip Step 2a / 2b /// - no -> try parent kernels /// - /// Step 2a: current_array.execute_parent(stack.top.parent_array) - /// child looks up at the suspended parent from ExecuteSlot + /// Step 2a: if stack.top exists: + /// parent = stack.top.parent_array + /// child = current_array + /// kernels[(parent.encoding_id(), child.encoding_id())] + /// .try_execute_parent(child, parent, stack.top.slot_idx) /// /// Step 2b: for child in current_array.children(): - /// child.execute_parent(current_array) - /// each child looks up at current_array + /// parent = current_array + /// kernels[(parent.encoding_id(), child.encoding_id())] + /// .try_execute_parent(child, parent, child.slot_idx) /// /// Step 3: match current_array.execute() /// ExecuteSlot(i, pred) -> push parent on stack, focus child `i` @@ -196,8 +200,17 @@ impl ArrayRef { // would be lost when we restore frame.parent_builder. if current_builder.is_none() && let Some(frame) = stack.last() - && let Some(result) = - current_array.execute_parent(&frame.parent_array, frame.slot_idx, ctx)? + && let Some(result) = { + let tmp_session = ctx.session().clone(); + let kernels = tmp_session.get_opt::(); + execute_parent_for_child( + &frame.parent_array, + ¤t_array, + frame.slot_idx, + kernels, + ctx, + )? + } { ctx.log(format_args!( "execute_parent (stack) rewrote {} -> {}", @@ -548,13 +561,13 @@ fn execute_parent_for_child( kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) { for plugin in plugins.as_ref() { - if let Some(result) = plugin(child, parent, slot_idx, ctx)? { + if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { return Ok(Some(result)); } } } - child.execute_parent(parent, slot_idx, ctx) + Ok(None) } /// Try execute_parent on each occupied slot of the array. diff --git a/vortex-array/src/kernel.rs b/vortex-array/src/kernel.rs index f5b75471437..dff9769f139 100644 --- a/vortex-array/src/kernel.rs +++ b/vortex-array/src/kernel.rs @@ -10,12 +10,10 @@ //! run ends rather than decoding the entire array and slicing the result. //! //! Encodings declare their parent kernels by implementing [`ExecuteParentKernel`] and -//! registering them in a [`ParentKernelSet`]. Each kernel specifies which parent types it -//! handles via a [`Matcher`]. +//! registering them with the session's optimizer-kernel registry. Each kernel specifies which +//! parent types it handles via a [`Matcher`]. -use std::any::type_name; use std::fmt::Debug; -use std::marker::PhantomData; use vortex_error::VortexResult; @@ -25,57 +23,6 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::matcher::Matcher; -/// A collection of [`ExecuteParentKernel`]s registered for a specific child encoding. -/// -/// During execution, the scheduler iterates over each child's `ParentKernelSet` looking for -/// 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], -} - -impl ParentKernelSet { - /// 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 } - } - - /// Lift the given rule into a dynamic trait object. - pub const fn lift>( - kernel: &'static K, - ) -> &'static dyn DynParentKernel { - // Assert that self is zero-sized - const { - assert!( - !(size_of::() != 0), - "Rule must be zero-sized to be lifted" - ); - } - unsafe { &*(kernel as *const K as *const ParentKernelAdapter) } - } - - /// Evaluate the parent kernels on the given child and parent arrays. - pub fn execute( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - for kernel in self.kernels.iter() { - if !kernel.matches(parent) { - continue; - } - if let Some(reduced) = kernel.execute_parent(child, parent, child_idx, ctx)? { - return Ok(Some(reduced)); - } - } - Ok(None) - } -} - /// A kernel that allows a child encoding `V` to execute its parent array in a fused manner. /// /// This is the typed trait that encoding authors implement. The associated `Parent` type @@ -99,55 +46,3 @@ pub trait ExecuteParentKernel: Debug + Send + Sync + 'static { ctx: &mut ExecutionCtx, ) -> VortexResult>; } - -/// Type-erased version of [`ExecuteParentKernel`] used for dynamic dispatch within -/// [`ParentKernelSet`]. -pub trait DynParentKernel: Send + Sync { - /// Returns `true` if this kernel's parent [`Matcher`] matches the given parent array. - fn matches(&self, parent: &ArrayRef) -> bool; - - /// Attempt to execute the parent array fused with the child array. - fn execute_parent( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; -} - -/// Bridges a concrete [`ExecuteParentKernel`] to the type-erased [`DynParentKernel`] -/// trait. Created by [`ParentKernelSet::lift`]. -pub struct ParentKernelAdapter { - kernel: K, - _phantom: PhantomData, -} - -impl> Debug for ParentKernelAdapter { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ParentKernelAdapter") - .field("parent", &type_name::()) - .field("kernel", &self.kernel) - .finish() - } -} - -impl> DynParentKernel for ParentKernelAdapter { - fn matches(&self, parent: &ArrayRef) -> bool { - K::Parent::matches(parent) - } - - fn execute_parent( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(parent_view) = K::Parent::try_match(parent) else { - return Ok(None); - }; - self.kernel - .execute_parent(child, parent_view, child_idx, ctx) - } -} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index b5ff5710489..5b3a70e6aa5 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -85,6 +85,11 @@ pub mod flatbuffers { pub use vortex_flatbuffers::array::*; } +/// Register vortex-array's built-in session-scoped kernels. +pub fn initialize(session: &VortexSession) { + arrays::initialize(session); +} + /// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session /// variables: arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, /// Arrow conversion, and memory. @@ -93,7 +98,7 @@ pub mod flatbuffers { /// additional encodings or kernels into it without affecting any other session. This does not /// register file, layout, or runtime state — those live in higher-level crates. pub fn array_session() -> VortexSession { - VortexSession::builder() + let session = VortexSession::builder() .with::() .with::() .with::() @@ -102,7 +107,9 @@ pub fn array_session() -> VortexSession { .with::() .with::() .with::() - .build() + .build(); + initialize(&session); + session } // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index b94393a49b3..357a57b4cd3 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -6,10 +6,9 @@ //! [`ArrayKernels`] stores function pointers that participate in array optimization and execution //! without adding rules or kernels to an encoding vtable. The optimizer consults it for //! parent-reduce rewrites before the child encoding's static `PARENT_RULES`, and the executor -//! consults it for parent execution before the child encoding's static parent kernels. A -//! registered function can therefore add support for an extension encoding or take precedence over -//! a built-in rule or kernel. When several functions are registered for the same key and kind, -//! they are tried in registration order until one applies. +//! consults it for parent execution. A registered function can therefore add support for an +//! extension encoding or take precedence over a built-in rule. When several functions are +//! registered for the same key and kind, they are tried in registration order until one applies. //! //! Kernel entries are addressed by `(outer_id, child_id)`. For parent-reduce and execute-parent //! kernels, `outer_id` is the id returned by the parent array's `encoding_id()` and `child_id` is @@ -25,6 +24,7 @@ use std::any::Any; use std::borrow::Borrow; +use std::fmt::Debug; use std::hash::BuildHasher; use std::sync::Arc; use std::sync::LazyLock; @@ -40,8 +40,9 @@ use crate::ExecutionCtx; use crate::arc_swap_map::ArcSwapMap; use crate::array::VTable; use crate::arrays::Struct; -use crate::arrays::struct_::compute::cast::struct_cast_execute_parent; use crate::arrays::struct_::compute::rules::struct_cast_reduce_parent; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -90,6 +91,64 @@ pub type ExecuteParentFn = fn( ctx: &mut ExecutionCtx, ) -> VortexResult>; +/// Type-erased execute-parent kernel stored in the session registry. +pub trait DynExecuteParentKernel: Debug + Send + Sync + 'static { + /// Attempt to execute the parent array fused with the child array. + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +type ExecuteParentKernelRef = Arc; + +#[derive(Debug)] +struct ExecuteParentFnKernel(ExecuteParentFn); + +impl DynExecuteParentKernel for ExecuteParentFnKernel { + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.0(child, parent, child_idx, ctx) + } +} + +#[derive(Debug)] +struct RegisteredExecuteParentKernel { + _child: V, + kernel: K, +} + +impl DynExecuteParentKernel for RegisteredExecuteParentKernel +where + V: VTable, + K: ExecuteParentKernel, +{ + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(child) = child.as_opt::() else { + return Ok(None); + }; + let Some(parent) = K::Parent::try_match(parent) else { + return Ok(None); + }; + + self.kernel.execute_parent(child, parent, child_idx, ctx) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] #[repr(transparent)] struct ExecuteParentFnId(u64); @@ -113,14 +172,13 @@ impl Borrow for ExecuteParentFnId { #[derive(Clone, Debug)] pub struct ArrayKernels { reduce_parent: ArcSwapMap>, - execute_parent: ArcSwapMap>, + execute_parent: ArcSwapMap>, } impl Default for ArrayKernels { fn default() -> ArrayKernels { let this = Self::empty(); this.register_builtin_reduce_parent(); - this.register_builtin_execute_parent(); this } } @@ -142,14 +200,6 @@ impl ArrayKernels { ); } - fn register_builtin_execute_parent(&self) { - self.register_execute_parent( - Cast.id(), - Struct.id(), - &[struct_cast_execute_parent as ExecuteParentFn], - ); - } - /// Register [`ReduceParentFn`]s for `(parent, child)`. /// /// The optimizer invokes these functions in registration order when it sees a parent with @@ -177,20 +227,48 @@ impl ArrayKernels { /// /// The executor invokes these functions in registration order when it sees a parent with /// encoding id `parent` holding a child with encoding id `child` during a parent execution - /// step, before trying the child encoding's static parent kernels. + /// step. /// /// If functions have already been registered for the same pair, these functions are appended /// after them. pub fn register_execute_parent(&self, parent: Id, child: Id, fns: &[ExecuteParentFn]) { + let kernels: Vec = fns + .iter() + .map(|f| Arc::new(ExecuteParentFnKernel(*f)) as ExecuteParentKernelRef) + .collect(); self.execute_parent - .extend(hash_fn_id(parent, child).into(), fns); + .extend(hash_fn_id(parent, child).into(), kernels.as_slice()); + } + + /// Register a typed [`ExecuteParentKernel`] for `(parent, child.id())`. + /// + /// The executor invokes registered kernels in registration order before falling through to + /// later registered kernels for the same key. `parent` is usually the parent array's encoding + /// id. For `ScalarFnArray`, it is the scalar function id, for example `Cast.id()`. + pub fn register_execute_parent_kernel(&self, parent: Id, child: V, kernel: K) + where + V: VTable, + K: ExecuteParentKernel, + { + let child_id = child.id(); + self.execute_parent.push( + hash_fn_id(parent, child_id).into(), + Arc::new(RegisteredExecuteParentKernel { + _child: child, + kernel, + }) as ExecuteParentKernelRef, + ); } /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`. /// /// Returns an owned [`Arc`] so the session-variable borrow can be dropped before invoking the /// functions. - pub fn find_execute_parent(&self, parent: Id, child: Id) -> Option> { + pub fn find_execute_parent( + &self, + parent: Id, + child: Id, + ) -> Option> { self.execute_parent.get(&hash_fn_id(parent, child)) } } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 7a131c807a3..dcbf7a6ddb3 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -108,17 +108,11 @@ pub use footer::*; pub use forever_constant::*; pub use open::*; pub use strategy::*; -use vortex_array::arrays::Dict; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; -use vortex_bytebool::ByteBool; -use vortex_fsst::FSST; -#[cfg(feature = "unstable_encodings")] -use vortex_onpair::OnPair; use vortex_pco::Pco; use vortex_session::VortexSession; -use vortex_zigzag::ZigZag; pub use writer::*; /// The current version of the Vortex file format @@ -158,15 +152,15 @@ mod forever_constant { /// NOTE: this function will be changed in the future to encapsulate logic for using different /// Vortex "Editions" that may support different sets of encodings. pub fn register_default_encodings(session: &VortexSession) { + vortex_bytebool::initialize(session); + vortex_fsst::initialize(session); + #[cfg(feature = "unstable_encodings")] + vortex_onpair::initialize(session); + vortex_zigzag::initialize(session); + { let arrays = session.arrays(); - arrays.register(ByteBool); - arrays.register(Dict); - arrays.register(FSST); - #[cfg(feature = "unstable_encodings")] - arrays.register(OnPair); arrays.register(Pco); - arrays.register(ZigZag); #[cfg(feature = "zstd")] arrays.register(vortex_zstd::Zstd); #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] @@ -176,8 +170,6 @@ pub fn register_default_encodings(session: &VortexSession) { } } - // Eventually all encodings crates should expose an initialize function. For now it's only - // a few of them. vortex_alp::initialize(session); vortex_datetime_parts::initialize(session); vortex_decimal_byte_parts::initialize(session); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index f0f94613c07..7f4e72630d8 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -176,6 +176,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); + vortex_array::initialize(&session); #[cfg(feature = "files")] let session = { From c52b6d5b219372275c6836c3c2e9e87a8541d4db Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 16:16:18 -0400 Subject: [PATCH 02/19] Fix session kernel benchmark setup Initialize direct encoding benchmark sessions with their crate-level kernel registrations so they exercise the session execute-parent path instead of fallback materialization. Cache the ArrayKernels handle in ExecutionCtx to avoid repeated session lookups while trying execute-parent kernels. Signed-off-by: "Nicholas Gates" Signed-off-by: Nicholas Gates --- encodings/alp/benches/alp_compress.rs | 6 ++++- .../experimental/onpair/benches/decode.rs | 6 ++++- .../fastlanes/benches/bitpack_compare.rs | 6 ++++- .../benches/bitpack_compare_sweep.rs | 6 ++++- .../fastlanes/benches/bitpacking_take.rs | 6 ++++- .../fastlanes/benches/canonicalize_bench.rs | 6 ++++- encodings/fastlanes/benches/cast_bitpacked.rs | 6 ++++- .../fastlanes/benches/compute_between.rs | 6 ++++- .../fsst/benches/chunked_dict_fsst_builder.rs | 6 ++++- encodings/fsst/benches/fsst_compress.rs | 6 ++++- encodings/fsst/benches/fsst_like.rs | 6 ++++- encodings/fsst/benches/fsst_url_compare.rs | 6 ++++- encodings/runend/benches/run_end_compress.rs | 6 ++++- encodings/runend/benches/run_end_decode.rs | 6 ++++- .../runend/benches/run_end_null_count.rs | 6 ++++- encodings/sparse/benches/sparse_canonical.rs | 6 ++++- vortex-array/src/executor.rs | 22 +++++++++++-------- 17 files changed, 93 insertions(+), 25 deletions(-) diff --git a/encodings/alp/benches/alp_compress.rs b/encodings/alp/benches/alp_compress.rs index f70f69c4177..02cc30e44a6 100644 --- a/encodings/alp/benches/alp_compress.rs +++ b/encodings/alp/benches/alp_compress.rs @@ -50,7 +50,11 @@ const BENCH_ARGS: &[(usize, f64, f64)] = &[ (10_000, 0.1, 1.0), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_alp::initialize(&session); + session +}); #[divan::bench(types = [f32, f64], args = BENCH_ARGS)] fn compress_alp(bencher: Bencher, args: (usize, f64, f64)) { diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index d8623018d8c..1ee96652891 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -79,7 +79,11 @@ impl DecodeInputs { use vortex_onpair::onpair_compress; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_onpair::initialize(&session); + session +}); #[derive(Copy, Clone, Debug)] enum Shape { diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index e1e8a0e365b..739fcd72dc6 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -35,7 +35,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); const LENS: &[usize] = &[1024, 64 * 1024]; const BIT_WIDTHS: &[u8] = &[4, 16]; diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index 037581d016b..ec7cf9b6892 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -39,7 +39,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); /// Number of elements per benchmarked array (64 full FastLanes blocks). const LEN: usize = 64 * 1024; diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 4d0a915aa99..eb072017ae3 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -26,7 +26,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); #[divan::bench] fn take_10_stratified(bencher: Bencher) { diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index deea4de0032..30ee0756960 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -34,7 +34,11 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (10000, 1000, 0.00), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); #[cfg(not(codspeed))] #[divan::bench(args = BENCH_ARGS)] diff --git a/encodings/fastlanes/benches/cast_bitpacked.rs b/encodings/fastlanes/benches/cast_bitpacked.rs index 9552d61c632..baeff20f5e1 100644 --- a/encodings/fastlanes/benches/cast_bitpacked.rs +++ b/encodings/fastlanes/benches/cast_bitpacked.rs @@ -38,7 +38,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); const U32: DType = DType::Primitive(PType::U32, Nullability::NonNullable); diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index 73af3cc07d7..2b295157e68 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -24,7 +24,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); fn generate_primitive_array( rng: &mut StdRng, diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index 2ea5cfb8f5d..8c2c6c2e388 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -28,7 +28,11 @@ const BENCH_ARGS: &[(usize, usize, usize)] = &[ (1000, 1000, 100), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); fn make_dict_fsst_chunks( len: usize, diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index cf648f68c48..6aad25de874 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -31,7 +31,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); // [(string_count, avg_len, unique_chars)] const BENCH_ARGS: &[(usize, usize, u8)] = &[ diff --git a/encodings/fsst/benches/fsst_like.rs b/encodings/fsst/benches/fsst_like.rs index 58a49af5582..414c1b7afaf 100644 --- a/encodings/fsst/benches/fsst_like.rs +++ b/encodings/fsst/benches/fsst_like.rs @@ -29,7 +29,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); const N: usize = NUM_STRINGS; diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index 00fdd2aa6f1..2a18236e6be 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -30,7 +30,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); const NUM_URLS: usize = NUM_STRINGS; diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index 9bb1b35ab94..cd241ad42c6 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -23,7 +23,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 4), diff --git a/encodings/runend/benches/run_end_decode.rs b/encodings/runend/benches/run_end_decode.rs index a2915c0b8b7..9256a4b3a53 100644 --- a/encodings/runend/benches/run_end_decode.rs +++ b/encodings/runend/benches/run_end_decode.rs @@ -20,7 +20,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); /// Distribution types for bool benchmarks #[derive(Clone, Copy)] diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index e247ad52208..f214889daae 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -49,7 +49,11 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (100_000, 1024, 0.5), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); #[divan::bench(args = BENCH_ARGS)] fn null_count_run_end(bencher: Bencher, (n, run_step, valid_density): (usize, usize, f64)) { diff --git a/encodings/sparse/benches/sparse_canonical.rs b/encodings/sparse/benches/sparse_canonical.rs index 8179bec369f..27f7340fb91 100644 --- a/encodings/sparse/benches/sparse_canonical.rs +++ b/encodings/sparse/benches/sparse_canonical.rs @@ -27,7 +27,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_sparse::initialize(&session); + session +}); const LIST_ARGS: &[(usize, usize, usize)] = &[ // len, patch_stride, list_size diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index c2f5eda9424..58b4ce8176b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -201,13 +201,12 @@ impl ArrayRef { if current_builder.is_none() && let Some(frame) = stack.last() && let Some(result) = { - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); + let kernels = ctx.array_kernels(); execute_parent_for_child( &frame.parent_array, ¤t_array, frame.slot_idx, - kernels, + kernels.as_ref(), ctx, )? } @@ -316,6 +315,7 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, + array_kernels: Option, #[cfg(debug_assertions)] id: usize, #[cfg(debug_assertions)] @@ -325,8 +325,10 @@ pub struct ExecutionCtx { impl ExecutionCtx { /// Create a new execution context with the given session. pub fn new(session: VortexSession) -> Self { + let array_kernels = session.get_opt::().cloned(); Self { session, + array_kernels, #[cfg(debug_assertions)] id: { static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0); @@ -342,6 +344,10 @@ impl ExecutionCtx { &self.session } + fn array_kernels(&self) -> Option { + self.array_kernels.clone() + } + /// Get the session-scoped host allocator for this execution context. pub fn allocator(&self) -> HostAllocatorRef { self.session.allocator() @@ -435,13 +441,12 @@ impl Executable for ArrayRef { } } - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); + let kernels = ctx.array_kernels(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(&array, child, slot_idx, kernels, ctx)? + execute_parent_for_child(&array, child, slot_idx, kernels.as_ref(), ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -572,13 +577,12 @@ fn execute_parent_for_child( /// Try execute_parent on each occupied slot of the array. fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); + let kernels = ctx.array_kernels(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(array, child, slot_idx, kernels, ctx)? + execute_parent_for_child(array, child, slot_idx, kernels.as_ref(), ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", From 6fdb9326281a5465e44f86a7d6cf250182a70289 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 16:42:28 -0400 Subject: [PATCH 03/19] Fix initialized sessions in encoding tests Signed-off-by: Nicholas Gates --- encodings/parquet-variant/src/kernel.rs | 13 +++++++++++-- encodings/runend/src/compute/filter.rs | 17 +++++++++++++---- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 424cb9d1d78..5624437d461 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -251,10 +251,17 @@ mod tests { use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; + fn session() -> VortexSession { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + } + fn make_unshredded_array() -> VortexResult { let mut builder = VariantArrayBuilder::new(4); builder.append_variant(PqVariant::from(42i32)); @@ -375,9 +382,10 @@ mod tests { dtype: Option, ) -> VortexResult { let expr = variant_get(root(), parse_path(path)?, dtype); + let session = session(); array .apply(&expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut session.create_execution_ctx()) } macro_rules! assert_rows_eq { @@ -826,7 +834,8 @@ mod tests { array: &ArrayRef, expected: impl IntoIterator>, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let session = session(); + let mut ctx = session.create_execution_ctx(); let executed = array.clone().execute::(&mut ctx)?; let typed_value = executed .as_::() diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 60644e2bced..609400d9e25 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -116,20 +116,27 @@ fn filter_run_end_primitive + AsPrimitiv #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_error::VortexResult; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; + fn session() -> VortexSession { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + } + fn ree_array() -> RunEndArray { + let session = session(); RunEnd::encode( PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut session.create_execution_ctx(), ) .unwrap() } @@ -139,7 +146,8 @@ mod tests { let arr = ree_array().slice(2..7)?; let filtered = arr.filter(Mask::from_iter([true, false, false, true, true]))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let session = session(); + let mut ctx = session.create_execution_ctx(); assert_arrays_eq!( filtered, RunEnd::new( @@ -157,7 +165,8 @@ mod tests { /// Filter unwrap one layer at a time so RunEnd's FilterKernel can fire. #[test] fn filter_sliced_run_end_preserves_encoding() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let session = session(); + let mut ctx = session.create_execution_ctx(); // 4 runs of 32 each = 128 rows. Large enough that FilterKernel takes // the run-preserving path (true_count >= 25). From 644e00d9862aaa1e01f38121b917f8e106279fa9 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 17:08:27 -0400 Subject: [PATCH 04/19] Initialize benchmark sessions after rebase Signed-off-by: Nicholas Gates --- encodings/runend/benches/run_end_take.rs | 8 +++++--- vortex/benches/common_encoding_tree_throughput.rs | 3 ++- vortex/benches/single_encoding_throughput.rs | 3 ++- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/encodings/runend/benches/run_end_take.rs b/encodings/runend/benches/run_end_take.rs index 5c5c75a712a..424590c3450 100644 --- a/encodings/runend/benches/run_end_take.rs +++ b/encodings/runend/benches/run_end_take.rs @@ -15,7 +15,6 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_runend::RunEnd; @@ -25,8 +24,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); #[derive(Clone, Copy)] enum IndexPattern { diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 69addd25a69..cefc9e10723 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -13,6 +13,7 @@ use divan::counter::BytesCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; +use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; @@ -52,7 +53,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(VortexSession::default); const NUM_VALUES: u64 = 100_000; diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index f70ec9e69f1..2c2db300d82 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -14,6 +14,7 @@ use rand::RngExt; use rand::SeedableRng; use rand::prelude::IndexedRandom; use rand::rngs::StdRng; +use vortex::VortexSessionDefault; use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; @@ -45,7 +46,7 @@ use vortex_session::VortexSession; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(VortexSession::default); fn main() { divan::main(); From afb594f153f1cbc4633fcbb4a06f12c17b582aa0 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 17:55:54 -0400 Subject: [PATCH 05/19] Cache execute-parent kernel lookups Signed-off-by: Nicholas Gates --- vortex-array/src/executor.rs | 70 ++++++++++++++++++++------- vortex-array/src/optimizer/kernels.rs | 15 +++++- 2 files changed, 67 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 58b4ce8176b..3df3f612bf9 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -15,12 +15,14 @@ use std::env::VarError; use std::fmt; use std::fmt::Display; +use std::sync::Arc; use std::sync::LazyLock; #[cfg(debug_assertions)] use std::sync::atomic::AtomicUsize; #[cfg(debug_assertions)] use std::sync::atomic::Ordering; +use smallvec::SmallVec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -41,6 +43,8 @@ use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ExecuteParentKernelRef; +use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; @@ -164,6 +168,8 @@ impl ArrayRef { let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); + let kernels = ctx.array_kernels.clone(); + let mut execute_parent_cache = ExecuteParentCache::default(); let max_iterations = max_iterations(); for _ in 0..max_iterations { @@ -201,12 +207,12 @@ impl ArrayRef { if current_builder.is_none() && let Some(frame) = stack.last() && let Some(result) = { - let kernels = ctx.array_kernels(); execute_parent_for_child( &frame.parent_array, ¤t_array, frame.slot_idx, kernels.as_ref(), + &mut execute_parent_cache, ctx, )? } @@ -223,7 +229,12 @@ impl ArrayRef { // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() - && let Some(rewritten) = try_execute_parent(¤t_array, ctx)? + && let Some(rewritten) = try_execute_parent( + ¤t_array, + kernels.as_ref(), + &mut execute_parent_cache, + ctx, + )? { ctx.log(format_args!( "execute_parent rewrote {} -> {}", @@ -344,10 +355,6 @@ impl ExecutionCtx { &self.session } - fn array_kernels(&self) -> Option { - self.array_kernels.clone() - } - /// Get the session-scoped host allocator for this execution context. pub fn allocator(&self) -> HostAllocatorRef { self.session.allocator() @@ -441,13 +448,19 @@ impl Executable for ArrayRef { } } - let kernels = ctx.array_kernels(); + let kernels = ctx.array_kernels.clone(); + let mut execute_parent_cache = ExecuteParentCache::default(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; - if let Some(executed_parent) = - execute_parent_for_child(&array, child, slot_idx, kernels.as_ref(), ctx)? - { + if let Some(executed_parent) = execute_parent_for_child( + &array, + child, + slot_idx, + kernels.as_ref(), + &mut execute_parent_cache, + ctx, + )? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, @@ -559,11 +572,11 @@ fn execute_parent_for_child( child: &ArrayRef, slot_idx: usize, kernels: Option<&ArrayKernels>, + cache: &mut ExecuteParentCache, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(kernels) = kernels - && let Some(plugins) = - kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) + if let Some(plugins) = + find_execute_parent(kernels, cache, parent.encoding_id(), child.encoding_id()) { for plugin in plugins.as_ref() { if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { @@ -576,13 +589,16 @@ fn execute_parent_for_child( } /// Try execute_parent on each occupied slot of the array. -fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - let kernels = ctx.array_kernels(); - +fn try_execute_parent( + array: &ArrayRef, + kernels: Option<&ArrayKernels>, + cache: &mut ExecuteParentCache, + ctx: &mut ExecutionCtx, +) -> VortexResult> { for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(array, child, slot_idx, kernels.as_ref(), ctx)? + execute_parent_for_child(array, child, slot_idx, kernels, cache, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -600,6 +616,26 @@ fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< Ok(None) } +type ExecuteParentCache = SmallVec<[(u64, Option>); 8]>; + +fn find_execute_parent( + kernels: Option<&ArrayKernels>, + cache: &mut ExecuteParentCache, + parent: ArrayId, + child: ArrayId, +) -> Option> { + let kernels = kernels?; + let key = execute_parent_key(parent, child); + + if let Some((_, cached)) = cache.iter().find(|(cached_key, _)| *cached_key == key) { + return cached.clone(); + } + + let found = kernels.find_execute_parent_by_key(key); + cache.push((key, found.clone())); + found +} + /// A predicate that determines when an array has reached a desired form during execution. pub type DonePredicate = fn(&ArrayRef) -> bool; diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 357a57b4cd3..3257d5800a5 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -103,7 +103,7 @@ pub trait DynExecuteParentKernel: Debug + Send + Sync + 'static { ) -> VortexResult>; } -type ExecuteParentKernelRef = Arc; +pub(crate) type ExecuteParentKernelRef = Arc; #[derive(Debug)] struct ExecuteParentFnKernel(ExecuteParentFn); @@ -271,12 +271,25 @@ impl ArrayKernels { ) -> Option> { self.execute_parent.get(&hash_fn_id(parent, child)) } + + /// Look up the execute-parent kernels registered for a pre-hashed key. + pub(crate) fn find_execute_parent_by_key( + &self, + key: u64, + ) -> Option> { + self.execute_parent.get(&key) + } } fn hash_fn_id(parent: Id, child: Id) -> u64 { FN_HASHER.hash_one((parent, child)) } +/// Return the registry key for execute-parent kernels registered for `(parent, child)`. +pub(crate) fn execute_parent_key(parent: Id, child: Id) -> u64 { + hash_fn_id(parent, child) +} + impl SessionVar for ArrayKernels { fn as_any(&self) -> &dyn Any { self From e22b3ef06b3c9be27c31c67d3b124e233a4c80d2 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 10:27:06 -0400 Subject: [PATCH 06/19] Fast-path canonical builder appends Signed-off-by: Nicholas Gates --- encodings/sparse/src/canonical.rs | 21 +++- vortex-array/src/arrays/bool/vtable/mod.rs | 15 +++ vortex-array/src/arrays/dict/vtable/mod.rs | 29 +++++ .../src/arrays/primitive/vtable/mod.rs | 18 +++ .../src/arrays/varbinview/vtable/mod.rs | 15 +++ vortex-array/src/builders/bool.rs | 27 ++--- vortex-array/src/builders/dict/mod.rs | 6 +- vortex-array/src/builders/fixed_size_list.rs | 12 +- vortex-array/src/builders/listview.rs | 12 +- vortex-array/src/builders/primitive.rs | 44 +++---- vortex-array/src/builders/varbinview.rs | 107 ++++++++++-------- 11 files changed, 216 insertions(+), 90 deletions(-) diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index dd84123851a..28af7c751e6 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -239,6 +239,7 @@ fn execute_sparse_lists_inner( &mut builder, fill_elements.as_ref(), sparse_idx - next_index, + ctx, ); if patch_valid { @@ -246,7 +247,7 @@ fn execute_sparse_lists_inner( .list_elements_at(patch_idx) .vortex_expect("list_elements_at"); builder - .append_array_as_list(&patch_list) + .append_array_as_list_ctx(&patch_list, ctx) .vortex_expect("Failed to append sparse value"); } else { builder.append_null(); @@ -255,7 +256,7 @@ fn execute_sparse_lists_inner( next_index = sparse_idx + 1; } - append_list_fill(&mut builder, fill_elements.as_ref(), len - next_index); + append_list_fill(&mut builder, fill_elements.as_ref(), len - next_index, ctx); builder.finish() } @@ -328,6 +329,7 @@ fn execute_sparse_fixed_size_list_inner( &mut builder, fill_elements.as_ref(), sparse_idx - next_index, + ctx, ); // Append the patch value, handling null patches by appending defaults. @@ -336,7 +338,7 @@ fn execute_sparse_fixed_size_list_inner( .fixed_size_list_elements_at(patch_idx) .vortex_expect("fixed_size_list_elements_at"); builder - .append_array_as_list(&patch_list) + .append_array_as_list_ctx(&patch_list, ctx) .vortex_expect("Failed to append sparse fixed-size-list value"); } else { builder.append_null(); @@ -346,7 +348,12 @@ fn execute_sparse_fixed_size_list_inner( } // Fill remaining positions after last patch. - append_fixed_size_list_fill(&mut builder, fill_elements.as_ref(), array_len - next_index); + append_fixed_size_list_fill( + &mut builder, + fill_elements.as_ref(), + array_len - next_index, + ctx, + ); builder.finish_into_fixed_size_list() } @@ -367,11 +374,12 @@ fn append_list_fill( builder: &mut ListViewBuilder, fill_elements: Option<&ArrayRef>, count: usize, + ctx: &mut ExecutionCtx, ) { if let Some(fill_elements) = fill_elements { for _ in 0..count { builder - .append_array_as_list(fill_elements) + .append_array_as_list_ctx(fill_elements, ctx) .vortex_expect("Failed to append sparse fill value"); } } else { @@ -383,11 +391,12 @@ fn append_fixed_size_list_fill( builder: &mut FixedSizeListBuilder, fill_elements: Option<&ArrayRef>, count: usize, + ctx: &mut ExecutionCtx, ) { if let Some(fill_elements) = fill_elements { for _ in 0..count { builder - .append_array_as_list(fill_elements) + .append_array_as_list_ctx(fill_elements, ctx) .vortex_expect("Failed to append sparse fixed-size-list fill value"); } } else { diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 264c3ab3a3b..eaa39462894 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -23,6 +23,8 @@ use crate::array::child_to_validity; use crate::arrays::bool::BoolData; use crate::arrays::bool::array::SLOT_NAMES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::BoolBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -173,6 +175,19 @@ impl VTable for Bool { SLOT_NAMES[idx].to_string() } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_bool_array(&array.into_owned(), ctx); + } + + builder.extend_from_array(array.as_ref()); + Ok(()) + } + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index bb9eb195132..cd2ef72a7ee 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -25,6 +25,7 @@ use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; use crate::EqMode; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -37,6 +38,7 @@ use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -199,6 +201,33 @@ impl VTable for Dict { )?)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if !array.is_empty() + && let (Some(codes), Some(values)) = ( + array.codes().as_opt::(), + array.values().as_opt::(), + ) + && !codes.validity()?.definitely_all_null() + { + let codes = codes.into_owned(); + let canonical = take_canonical(values, &codes, ctx)?.into_array(); + canonical.append_to_builder(builder, ctx)?; + return Ok(()); + } + + let canonical = array + .array() + .clone() + .execute::(ctx)? + .into_array(); + canonical.append_to_builder(builder, ctx)?; + Ok(()) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 90f56ea1aa5..d9d778c34bb 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -14,8 +14,11 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::primitive::PrimitiveData; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::PrimitiveBuilder; use crate::dtype::DType; use crate::dtype::PType; +use crate::match_each_native_ptype; use crate::serde::ArrayChildren; use crate::validity::Validity; mod kernel; @@ -187,6 +190,21 @@ impl VTable for Primitive { Ok(ExecutionResult::done(array)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match_each_native_ptype!(array.ptype(), |P| { + if let Some(builder) = builder.as_any_mut().downcast_mut::>() { + return builder.append_primitive_array(&array.into_owned(), ctx); + } + }); + + builder.extend_from_array(array.as_ref()); + Ok(()) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index d6d6b98aa7e..a9ffe5a5859 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -28,6 +28,8 @@ use crate::arrays::varbinview::array::NUM_SLOTS; use crate::arrays::varbinview::array::SLOT_NAMES; use crate::arrays::varbinview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -220,6 +222,19 @@ impl VTable for VarBinView { PARENT_RULES.evaluate(array, parent, child_idx) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_varbinview_array(&array.into_owned(), ctx); + } + + builder.extend_from_array(array.as_ref()); + Ok(()) + } + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index e829a58f6ec..221ceb75538 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -71,6 +72,17 @@ impl BoolBuilder { self.nulls.finish_with_nullability(self.dtype.nullability()), ) } + + pub(crate) fn append_bool_array( + &mut self, + array: &BoolArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + self.inner.append_buffer(&array.to_bit_buffer()); + self.nulls + .append_validity_mask(&BoolArrayExt::validity(array).execute_mask(array.len(), ctx)?); + Ok(()) + } } impl ArrayBuilder for BoolBuilder { @@ -118,19 +130,8 @@ impl ArrayBuilder for BoolBuilder { unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { #[expect(deprecated)] let bool_array = array.to_bool(); - - self.inner.append_buffer(&bool_array.to_bit_buffer()); - self.nulls.append_validity_mask( - &bool_array - .as_ref() - .validity() - .vortex_expect("validity_mask") - .execute_mask( - bool_array.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("Failed to compute validity mask"), - ); + self.append_bool_array(&bool_array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append bool array"); } fn reserve_exact(&mut self, additional: usize) { diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index c64e436bb45..a5a5bea2c9f 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -69,7 +69,11 @@ pub fn dict_encode_with_constraints( ) -> VortexResult { let mut encoder = dict_encoder(array, constraints); let encoded = encoder.encode(array); - let codes = encoded.execute::(ctx)?.narrow(ctx)?; + let codes = if let Some(codes) = encoded.as_opt::() { + codes.narrow(ctx)? + } else { + encoded.execute::(ctx)?.narrow(ctx)? + }; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code // since we build the dictionary from the codes we observe during encoding diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7a3c3bb068a..463e1ea8c04 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -12,6 +12,7 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -84,6 +85,15 @@ impl FixedSizeListBuilder { /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + self.append_array_as_list_ctx(array, &mut LEGACY_SESSION.create_execution_ctx()) + } + + /// Appends an array as a single non-null list entry using the provided execution context. + pub fn append_array_as_list_ctx( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -97,7 +107,7 @@ impl FixedSizeListBuilder { self.list_size() ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); Ok(()) diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 4c350f1888d..4e396186f9b 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -20,6 +20,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::ExecutionCtx; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::IntoArray; @@ -118,6 +119,15 @@ impl ListViewBuilder { /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + self.append_array_as_list_ctx(array, &mut LEGACY_SESSION.create_execution_ctx()) + } + + /// Appends an array as a single non-null list entry using the provided execution context. + pub fn append_array_as_list_ctx( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -135,7 +145,7 @@ impl ListViewBuilder { "appending this list would cause an offset overflow" ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index c4b5abfb083..3738cb65682 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -131,6 +132,28 @@ impl PrimitiveBuilder { self.values.extend(iter); self.nulls.append_validity_mask(mask); } + + pub(crate) fn append_primitive_array( + &mut self, + array: &PrimitiveArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + debug_assert_eq!( + array.ptype(), + T::PTYPE, + "Cannot append primitive array with different ptype" + ); + + self.values.extend_from_slice(array.as_slice::()); + self.nulls.append_validity_mask( + &array + .as_ref() + .validity() + .vortex_expect("validity_mask") + .execute_mask(array.as_ref().len(), ctx)?, + ); + Ok(()) + } } impl ArrayBuilder for PrimitiveBuilder { @@ -181,25 +204,8 @@ impl ArrayBuilder for PrimitiveBuilder { #[expect(deprecated)] let array = array.to_primitive(); - // This should be checked in `extend_from_array` but we can check it again. - debug_assert_eq!( - array.ptype(), - T::PTYPE, - "Cannot extend from array with different ptype" - ); - - self.values.extend_from_slice(array.as_slice::()); - self.nulls.append_validity_mask( - &array - .as_ref() - .validity() - .vortex_expect("validity_mask") - .execute_mask( - array.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("Failed to compute validity mask"), - ); + self.append_primitive_array(&array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append primitive array"); } fn reserve_exact(&mut self, additional: usize) { diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 392bef59997..bfb56a27670 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -19,10 +19,12 @@ use vortex_utils::aliases::hash_map::Entry; use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::VarBinViewArrayExt; use crate::arrays::varbinview::build_views::BinaryView; use crate::arrays::varbinview::compact::BufferUtilization; use crate::builders::ArrayBuilder; @@ -239,6 +241,60 @@ impl VarBinViewBuilder { fn push_only_validity_mask(&mut self, validity_mask: &Mask) { self.nulls.append_validity_mask(validity_mask); } + + pub(crate) fn append_varbinview_array( + &mut self, + array: &VarBinViewArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + self.flush_in_progress(); + + let mask = array.varbinview_validity().execute_mask(array.len(), ctx)?; + + self.push_only_validity_mask(&mask); + + let view_adjustment = + self.completed + .extend_from_compaction(BuffersWithOffsets::from_array( + array, + self.compaction_threshold, + )); + + match view_adjustment { + ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted( + array + .views() + .iter() + .map(|view| adjustment.adjust_view(view)), + ), + ViewAdjustment::Rewriting(adjustment) => match mask { + Mask::AllTrue(_) => { + for (idx, &view) in array.views().iter().enumerate() { + let new_view = self.push_view(view, &adjustment, array, idx); + self.views_builder.push(new_view); + } + } + Mask::AllFalse(_) => { + self.views_builder + .push_n(BinaryView::empty_view(), array.len()); + } + Mask::Values(v) => { + for (idx, (&view, is_valid)) in + array.views().iter().zip(v.bit_buffer().iter()).enumerate() + { + let new_view = if !is_valid { + BinaryView::empty_view() + } else { + self.push_view(view, &adjustment, array, idx) + }; + self.views_builder.push(new_view); + } + } + }, + } + + Ok(()) + } } impl ArrayBuilder for VarBinViewBuilder { @@ -297,55 +353,8 @@ impl ArrayBuilder for VarBinViewBuilder { unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { #[expect(deprecated)] let array = array.to_varbinview(); - self.flush_in_progress(); - - let mask = array - .validity() - .vortex_expect("validity_mask") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("Failed to compute validity mask"); - - self.push_only_validity_mask(&mask); - - let view_adjustment = - self.completed - .extend_from_compaction(BuffersWithOffsets::from_array( - &array, - self.compaction_threshold, - )); - - match view_adjustment { - ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted( - array - .views() - .iter() - .map(|view| adjustment.adjust_view(view)), - ), - ViewAdjustment::Rewriting(adjustment) => match mask { - Mask::AllTrue(_) => { - for (idx, &view) in array.views().iter().enumerate() { - let new_view = self.push_view(view, &adjustment, &array, idx); - self.views_builder.push(new_view); - } - } - Mask::AllFalse(_) => { - self.views_builder - .push_n(BinaryView::empty_view(), array.len()); - } - Mask::Values(v) => { - for (idx, (&view, is_valid)) in - array.views().iter().zip(v.bit_buffer().iter()).enumerate() - { - let new_view = if !is_valid { - BinaryView::empty_view() - } else { - self.push_view(view, &adjustment, &array, idx) - }; - self.views_builder.push(new_view); - } - } - }, - } + self.append_varbinview_array(&array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append varbinview array"); } fn reserve_exact(&mut self, additional: usize) { From 10aa40611d6d413d02502cdb2d1d69bac7e1c3cc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 10:39:53 -0400 Subject: [PATCH 07/19] Thread execution context through list appends Signed-off-by: Nicholas Gates --- encodings/sparse/src/canonical.rs | 8 +++---- vortex-array/src/builders/dict/mod.rs | 6 +---- vortex-array/src/builders/fixed_size_list.rs | 24 ++++++++++++-------- vortex-array/src/builders/list.rs | 22 +++++++++++++----- vortex-array/src/builders/listview.rs | 20 ++++++++-------- 5 files changed, 45 insertions(+), 35 deletions(-) diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 28af7c751e6..17b4241f293 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -247,7 +247,7 @@ fn execute_sparse_lists_inner( .list_elements_at(patch_idx) .vortex_expect("list_elements_at"); builder - .append_array_as_list_ctx(&patch_list, ctx) + .append_array_as_list(&patch_list, ctx) .vortex_expect("Failed to append sparse value"); } else { builder.append_null(); @@ -338,7 +338,7 @@ fn execute_sparse_fixed_size_list_inner( .fixed_size_list_elements_at(patch_idx) .vortex_expect("fixed_size_list_elements_at"); builder - .append_array_as_list_ctx(&patch_list, ctx) + .append_array_as_list(&patch_list, ctx) .vortex_expect("Failed to append sparse fixed-size-list value"); } else { builder.append_null(); @@ -379,7 +379,7 @@ fn append_list_fill( if let Some(fill_elements) = fill_elements { for _ in 0..count { builder - .append_array_as_list_ctx(fill_elements, ctx) + .append_array_as_list(fill_elements, ctx) .vortex_expect("Failed to append sparse fill value"); } } else { @@ -396,7 +396,7 @@ fn append_fixed_size_list_fill( if let Some(fill_elements) = fill_elements { for _ in 0..count { builder - .append_array_as_list_ctx(fill_elements, ctx) + .append_array_as_list(fill_elements, ctx) .vortex_expect("Failed to append sparse fixed-size-list fill value"); } } else { diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index a5a5bea2c9f..c64e436bb45 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -69,11 +69,7 @@ pub fn dict_encode_with_constraints( ) -> VortexResult { let mut encoder = dict_encoder(array, constraints); let encoded = encoder.encode(array); - let codes = if let Some(codes) = encoded.as_opt::() { - codes.narrow(ctx)? - } else { - encoded.execute::(ctx)?.narrow(ctx)? - }; + let codes = encoded.execute::(ctx)?.narrow(ctx)?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code // since we build the dictionary from the codes we observe during encoding diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 463e1ea8c04..2e7aced758b 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -84,12 +84,7 @@ impl FixedSizeListBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { - self.append_array_as_list_ctx(array, &mut LEGACY_SESSION.create_execution_ctx()) - } - - /// Appends an array as a single non-null list entry using the provided execution context. - pub fn append_array_as_list_ctx( + pub fn append_array_as_list( &mut self, array: &ArrayRef, ctx: &mut ExecutionCtx, @@ -1007,12 +1002,13 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -1028,7 +1024,7 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5, 6].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Interleave with another list scalar. builder @@ -1058,11 +1054,19 @@ mod tests { let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); // Test length mismatch error. let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 10); let wrong_len_arr = buffer![1i32, 2].into_array(); - assert!(builder.append_array_as_list(&wrong_len_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_len_arr, &mut ctx) + .is_err() + ); } } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ac97230daa8..c23e35cccff 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -13,6 +13,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -95,7 +96,11 @@ impl ListBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] in nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + pub fn append_array_as_list( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -103,7 +108,7 @@ impl ListBuilder { self.element_dtype() ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( O::from_usize(self.elements_builder.len()) @@ -637,12 +642,13 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -658,11 +664,11 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Append an empty array as a single list entry (empty list). let arr3 = buffer![0i32; 0].into_array(); - builder.append_array_as_list(&arr3).unwrap(); + builder.append_array_as_list(&arr3, &mut ctx).unwrap(); // Interleave with another list scalar (empty list). builder @@ -687,6 +693,10 @@ mod tests { // Test dtype mismatch error. let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); } } diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 4e396186f9b..fef03662181 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -118,12 +118,7 @@ impl ListViewBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { - self.append_array_as_list_ctx(array, &mut LEGACY_SESSION.create_execution_ctx()) - } - - /// Appends an array as a single non-null list entry using the provided execution context. - pub fn append_array_as_list_ctx( + pub fn append_array_as_list( &mut self, array: &ArrayRef, ctx: &mut ExecutionCtx, @@ -783,12 +778,13 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let mut builder = ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -804,11 +800,11 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Append an empty array as a single list entry (empty list). let arr3 = buffer![0i32; 0].into_array(); - builder.append_array_as_list(&arr3).unwrap(); + builder.append_array_as_list(&arr3, &mut ctx).unwrap(); // Interleave with another list scalar. builder @@ -839,6 +835,10 @@ mod tests { // Test dtype mismatch error. let mut builder = ListViewBuilder::::with_capacity(dtype, NonNullable, 20, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); } } From 455ae129148035e03b99759026c262f68c7b8b11 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 10:48:48 -0400 Subject: [PATCH 08/19] Share encoding test sessions Signed-off-by: Nicholas Gates --- .../src/bitpacking/vtable/operations.rs | 9 +-------- encodings/fastlanes/src/lib.rs | 6 +----- encodings/runend/src/compute/filter.rs | 17 ++++------------- encodings/runend/src/lib.rs | 9 +++++++++ 4 files changed, 15 insertions(+), 26 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index a41d3dcebca..92ff63f4408 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -31,7 +31,6 @@ impl OperationsVTable for BitPacked { #[cfg(test)] mod test { use std::ops::Range; - use std::sync::LazyLock; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -52,18 +51,12 @@ mod test { use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; - use vortex_session::VortexSession; use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; - - static FASTLANES_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session - }); + use crate::test::SESSION as FASTLANES_SESSION; fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 84d87649028..458b5d7563f 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -140,7 +140,6 @@ mod test { use std::sync::LazyLock; use vortex_array::VortexSessionExecute; - use vortex_array::session::ArraySessionExt; use vortex_buffer::BitBufferMut; use vortex_session::VortexSession; @@ -148,10 +147,7 @@ mod test { pub static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); - session.arrays().register(BitPacked); - session.arrays().register(Delta); - session.arrays().register(FoR); - session.arrays().register(RLE); + initialize(&session); session }); diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 609400d9e25..83eb811e682 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -121,22 +121,15 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_error::VortexResult; use vortex_mask::Mask; - use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; - - fn session() -> VortexSession { - let session = vortex_array::array_session(); - crate::initialize(&session); - session - } + use crate::tests::SESSION; fn ree_array() -> RunEndArray { - let session = session(); RunEnd::encode( PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(), - &mut session.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() } @@ -146,8 +139,7 @@ mod tests { let arr = ree_array().slice(2..7)?; let filtered = arr.filter(Mask::from_iter([true, false, false, true, true]))?; - let session = session(); - let mut ctx = session.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( filtered, RunEnd::new( @@ -165,8 +157,7 @@ mod tests { /// Filter unwrap one layer at a time so RunEnd's FilterKernel can fire. #[test] fn filter_sliced_run_end_preserves_encoding() -> VortexResult<()> { - let session = session(); - let mut ctx = session.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 4 runs of 32 each = 128 rows. Large enough that FilterKernel takes // the run-preserving path (true_count >= 25). diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 846c46d05d4..8169f5d8dfc 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -60,12 +60,21 @@ pub fn initialize(session: &VortexSession) { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use prost::Message; use vortex_array::dtype::PType; use vortex_array::test_harness::check_metadata; + use vortex_session::VortexSession; use crate::RunEndMetadata; + pub static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[cfg_attr(miri, ignore)] #[test] fn test_runend_metadata() { From 677d25174d130329b5d2151a341743ad3cf3281a Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 11:31:43 -0400 Subject: [PATCH 09/19] Snapshot execute-parent kernels Signed-off-by: "Nicholas Gates" --- vortex-array/src/arc_swap_map.rs | 20 ++++++++++++++ vortex-array/src/executor.rs | 40 +++++++++++++++------------ vortex-array/src/optimizer/kernels.rs | 14 +++++----- 3 files changed, 49 insertions(+), 25 deletions(-) diff --git a/vortex-array/src/arc_swap_map.rs b/vortex-array/src/arc_swap_map.rs index e8308d6343c..73108ae5d57 100644 --- a/vortex-array/src/arc_swap_map.rs +++ b/vortex-array/src/arc_swap_map.rs @@ -55,6 +55,11 @@ impl Debug for ArcSwapMap { } impl ArcSwapMap { + /// Return the currently published map snapshot. + pub(crate) fn snapshot(&self) -> Arc> { + self.inner.load_full() + } + /// Read the current snapshot, passing it to `f`. /// /// Every lookup inside `f` observes the same snapshot, which matters when a @@ -162,6 +167,21 @@ mod tests { assert_eq!(map.read(|m| m.values().sum::()), 3); } + #[test] + fn snapshot_keeps_published_view() { + let map = ArcSwapMap::::default(); + map.insert(1, 10); + + let snapshot = map.snapshot(); + map.insert(1, 20); + map.insert(2, 30); + + assert_eq!(snapshot.get(&1), Some(&10)); + assert_eq!(snapshot.get(&2), None); + assert_eq!(map.get(&1), Some(20)); + assert_eq!(map.get(&2), Some(30)); + } + #[test] fn clone_shares_the_same_cell() { let map = ArcSwapMap::::default(); diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 3df3f612bf9..df6bee53a26 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -44,10 +44,14 @@ use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ExecuteParentKernelRef; +use crate::optimizer::kernels::ParentExecutionKernels; use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; +static EMPTY_PARENT_EXECUTION_KERNELS: LazyLock> = + LazyLock::new(|| Arc::new(ParentExecutionKernels::default())); + /// 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 2^22. pub(crate) fn max_iterations() -> usize { @@ -168,7 +172,8 @@ impl ArrayRef { let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); - let kernels = ctx.array_kernels.clone(); + let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); + let kernels = execute_parent_kernels.as_ref(); let mut execute_parent_cache = ExecuteParentCache::default(); let max_iterations = max_iterations(); @@ -211,7 +216,7 @@ impl ArrayRef { &frame.parent_array, ¤t_array, frame.slot_idx, - kernels.as_ref(), + kernels, &mut execute_parent_cache, ctx, )? @@ -229,12 +234,8 @@ impl ArrayRef { // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() - && let Some(rewritten) = try_execute_parent( - ¤t_array, - kernels.as_ref(), - &mut execute_parent_cache, - ctx, - )? + && let Some(rewritten) = + try_execute_parent(¤t_array, kernels, &mut execute_parent_cache, ctx)? { ctx.log(format_args!( "execute_parent rewrote {} -> {}", @@ -326,7 +327,7 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, - array_kernels: Option, + execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, #[cfg(debug_assertions)] @@ -336,10 +337,13 @@ pub struct ExecutionCtx { impl ExecutionCtx { /// Create a new execution context with the given session. pub fn new(session: VortexSession) -> Self { - let array_kernels = session.get_opt::().cloned(); + let execute_parent_kernels = session + .get_opt::() + .map(ArrayKernels::execute_parent_snapshot) + .unwrap_or_else(|| Arc::clone(&EMPTY_PARENT_EXECUTION_KERNELS)); Self { session, - array_kernels, + execute_parent_kernels, #[cfg(debug_assertions)] id: { static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0); @@ -448,7 +452,8 @@ impl Executable for ArrayRef { } } - let kernels = ctx.array_kernels.clone(); + let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); + let kernels = execute_parent_kernels.as_ref(); let mut execute_parent_cache = ExecuteParentCache::default(); for (slot_idx, slot) in array.slots().iter().enumerate() { @@ -457,7 +462,7 @@ impl Executable for ArrayRef { &array, child, slot_idx, - kernels.as_ref(), + kernels, &mut execute_parent_cache, ctx, )? { @@ -571,7 +576,7 @@ fn execute_parent_for_child( parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, - kernels: Option<&ArrayKernels>, + kernels: &ParentExecutionKernels, cache: &mut ExecuteParentCache, ctx: &mut ExecutionCtx, ) -> VortexResult> { @@ -591,7 +596,7 @@ fn execute_parent_for_child( /// Try execute_parent on each occupied slot of the array. fn try_execute_parent( array: &ArrayRef, - kernels: Option<&ArrayKernels>, + kernels: &ParentExecutionKernels, cache: &mut ExecuteParentCache, ctx: &mut ExecutionCtx, ) -> VortexResult> { @@ -619,19 +624,18 @@ fn try_execute_parent( type ExecuteParentCache = SmallVec<[(u64, Option>); 8]>; fn find_execute_parent( - kernels: Option<&ArrayKernels>, + kernels: &ParentExecutionKernels, cache: &mut ExecuteParentCache, parent: ArrayId, child: ArrayId, ) -> Option> { - let kernels = kernels?; let key = execute_parent_key(parent, child); if let Some((_, cached)) = cache.iter().find(|(cached_key, _)| *cached_key == key) { return cached.clone(); } - let found = kernels.find_execute_parent_by_key(key); + let found = kernels.get(&key).cloned(); cache.push((key, found.clone())); found } diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 3257d5800a5..f2f89498b97 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -34,6 +34,7 @@ use vortex_session::SessionExt; use vortex_session::SessionVar; use vortex_session::registry::Id; use vortex_utils::aliases::DefaultHashBuilder; +use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; use crate::ExecutionCtx; @@ -105,6 +106,8 @@ pub trait DynExecuteParentKernel: Debug + Send + Sync + 'static { pub(crate) type ExecuteParentKernelRef = Arc; +pub(crate) type ParentExecutionKernels = HashMap>; + #[derive(Debug)] struct ExecuteParentFnKernel(ExecuteParentFn); @@ -151,7 +154,7 @@ where #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] #[repr(transparent)] -struct ExecuteParentFnId(u64); +pub(crate) struct ExecuteParentFnId(u64); impl From for ExecuteParentFnId { fn from(id: u64) -> Self { @@ -272,12 +275,9 @@ impl ArrayKernels { self.execute_parent.get(&hash_fn_id(parent, child)) } - /// Look up the execute-parent kernels registered for a pre-hashed key. - pub(crate) fn find_execute_parent_by_key( - &self, - key: u64, - ) -> Option> { - self.execute_parent.get(&key) + /// Return the currently published execute-parent kernel snapshot. + pub(crate) fn execute_parent_snapshot(&self) -> Arc { + self.execute_parent.snapshot() } } From 6d65ddd094b1a64b24db3efbf99ad3ef30bda4b9 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 12:02:16 -0400 Subject: [PATCH 10/19] Reduce parent execution dispatch overhead Avoid the per-execution parent kernel cache now that ExecutionCtx already holds the session snapshot, and make patch index lookup use the primitive fast path directly. Signed-off-by: "Nicholas Gates" --- .../src/arrays/filter/execute/take/tests.rs | 24 ++++----- vortex-array/src/executor.rs | 52 +++---------------- vortex-array/src/patches.rs | 7 +-- 3 files changed, 20 insertions(+), 63 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index 96728f988a6..ecfebe46679 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -9,7 +9,9 @@ use crate::IntoArray; use crate::RecursiveCanonical; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; +use crate::arrays::Dict; use crate::arrays::DictArray; +use crate::arrays::Filter; use crate::arrays::FilterArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; @@ -17,11 +19,12 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::arrays::dict::TakeExecuteAdaptor; use crate::assert_arrays_eq; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::executor::ExecutionCtx; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; fn execute_parent( @@ -30,19 +33,12 @@ fn execute_parent( child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let kernels = ctx.session().kernels().clone(); - let Some(plugins) = kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) - else { - return Ok(None); - }; - - for plugin in plugins.as_ref() { - if let Some(result) = plugin.execute_parent(child, parent, child_idx, ctx)? { - return Ok(Some(result)); - } - } - - Ok(None) + TakeExecuteAdaptor(Filter).execute_parent( + child.as_::(), + parent.as_::(), + child_idx, + ctx, + ) } #[test] diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index df6bee53a26..f90400c3412 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -22,7 +22,6 @@ use std::sync::atomic::AtomicUsize; #[cfg(debug_assertions)] use std::sync::atomic::Ordering; -use smallvec::SmallVec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -43,15 +42,11 @@ use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernels; -use crate::optimizer::kernels::ExecuteParentKernelRef; use crate::optimizer::kernels::ParentExecutionKernels; use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; -static EMPTY_PARENT_EXECUTION_KERNELS: LazyLock> = - LazyLock::new(|| Arc::new(ParentExecutionKernels::default())); - /// 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 2^22. pub(crate) fn max_iterations() -> usize { @@ -174,7 +169,6 @@ impl ArrayRef { let mut stack: Vec = Vec::new(); let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); let kernels = execute_parent_kernels.as_ref(); - let mut execute_parent_cache = ExecuteParentCache::default(); let max_iterations = max_iterations(); for _ in 0..max_iterations { @@ -217,7 +211,6 @@ impl ArrayRef { ¤t_array, frame.slot_idx, kernels, - &mut execute_parent_cache, ctx, )? } @@ -234,8 +227,7 @@ impl ArrayRef { // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() - && let Some(rewritten) = - try_execute_parent(¤t_array, kernels, &mut execute_parent_cache, ctx)? + && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { ctx.log(format_args!( "execute_parent rewrote {} -> {}", @@ -340,7 +332,7 @@ impl ExecutionCtx { let execute_parent_kernels = session .get_opt::() .map(ArrayKernels::execute_parent_snapshot) - .unwrap_or_else(|| Arc::clone(&EMPTY_PARENT_EXECUTION_KERNELS)); + .unwrap_or_default(); Self { session, execute_parent_kernels, @@ -454,18 +446,12 @@ impl Executable for ArrayRef { let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); let kernels = execute_parent_kernels.as_ref(); - let mut execute_parent_cache = ExecuteParentCache::default(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; - if let Some(executed_parent) = execute_parent_for_child( - &array, - child, - slot_idx, - kernels, - &mut execute_parent_cache, - ctx, - )? { + if let Some(executed_parent) = + execute_parent_for_child(&array, child, slot_idx, kernels, ctx)? + { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, @@ -577,12 +563,10 @@ fn execute_parent_for_child( child: &ArrayRef, slot_idx: usize, kernels: &ParentExecutionKernels, - cache: &mut ExecuteParentCache, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(plugins) = - find_execute_parent(kernels, cache, parent.encoding_id(), child.encoding_id()) - { + let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); + if let Some(plugins) = kernels.get(&key) { for plugin in plugins.as_ref() { if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { return Ok(Some(result)); @@ -597,13 +581,12 @@ fn execute_parent_for_child( fn try_execute_parent( array: &ArrayRef, kernels: &ParentExecutionKernels, - cache: &mut ExecuteParentCache, ctx: &mut ExecutionCtx, ) -> VortexResult> { for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(array, child, slot_idx, kernels, cache, ctx)? + execute_parent_for_child(array, child, slot_idx, kernels, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -621,25 +604,6 @@ fn try_execute_parent( Ok(None) } -type ExecuteParentCache = SmallVec<[(u64, Option>); 8]>; - -fn find_execute_parent( - kernels: &ParentExecutionKernels, - cache: &mut ExecuteParentCache, - parent: ArrayId, - child: ArrayId, -) -> Option> { - let key = execute_parent_key(parent, child); - - if let Some((_, cached)) = cache.iter().find(|(cached_key, _)| *cached_key == key) { - return cached.clone(); - } - - let found = kernels.get(&key).cloned(); - cache.push((key, found.clone())); - found -} - /// A predicate that determines when an array has reached a desired form during execution. pub type DonePredicate = fn(&ArrayRef) -> bool; diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index f4ab1cb0f9a..76f1c8ad99c 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -24,9 +24,8 @@ use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; @@ -497,9 +496,7 @@ impl Patches { /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] /// with the insertion point if not found. fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { - if indices.is_canonical() { - #[expect(deprecated)] - let primitive = indices.to_primitive(); + if let Some(primitive) = indices.as_opt::() { match_each_integer_ptype!(primitive.ptype(), |T| { let Ok(needle) = T::try_from(needle) else { // If the needle is not of type T, then it cannot possibly be in this array. From 7a0eb77a5fcba68d07b29e0ce863f6bc8c7a863b Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 12:43:15 -0400 Subject: [PATCH 11/19] Register built-in kernels with ArraySession Store vortex-array's built-in kernel registry on ArraySession so sessions that install the built-in array encodings also get their matching execute-parent kernels. Signed-off-by: "Nicholas Gates" --- vortex-array/src/executor.rs | 6 ++--- vortex-array/src/lib.rs | 17 +++++++----- vortex-array/src/optimizer/kernels.rs | 24 +++++++++++++---- vortex-array/src/optimizer/mod.rs | 4 +-- vortex-array/src/session/mod.rs | 39 +++++++++++++++++++++++++++ vortex/src/lib.rs | 3 --- 6 files changed, 73 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index f90400c3412..600d07d6d52 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -41,7 +41,7 @@ use crate::matcher::Matcher; use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; -use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::kernels::ParentExecutionKernels; use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; @@ -330,8 +330,8 @@ impl ExecutionCtx { /// Create a new execution context with the given session. pub fn new(session: VortexSession) -> Self { let execute_parent_kernels = session - .get_opt::() - .map(ArrayKernels::execute_parent_snapshot) + .kernels_opt() + .map(|kernels| kernels.execute_parent_snapshot()) .unwrap_or_default(); Self { session, diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 5b3a70e6aa5..2c0ce752329 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -85,9 +85,15 @@ pub mod flatbuffers { pub use vortex_flatbuffers::array::*; } -/// Register vortex-array's built-in session-scoped kernels. +/// Register vortex-array's built-in session-scoped kernels into a standalone +/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. +/// +/// Sessions that use [`ArraySession`] already receive vortex-array's built-in kernels from +/// [`ArraySession::default`]. pub fn initialize(session: &VortexSession) { - arrays::initialize(session); + if session.get_opt::().is_some() { + arrays::initialize(session); + } } /// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session @@ -98,18 +104,15 @@ pub fn initialize(session: &VortexSession) { /// additional encodings or kernels into it without affecting any other session. This does not /// register file, layout, or runtime state — those live in higher-level crates. pub fn array_session() -> VortexSession { - let session = VortexSession::builder() + VortexSession::builder() .with::() .with::() .with::() .with::() - .with::() .with::() .with::() .with::() - .build(); - initialize(&session); - session + .build() } // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index f2f89498b97..c13544c2b42 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -18,9 +18,10 @@ //! Because registered functions have different signatures for each kernel kind, the registry //! maintains one storage map per function type rather than a single type-erased map. //! -//! Sessions created by the top-level `vortex` crate install the default registry. Other sessions -//! can add it with [`VortexSession::with`](vortex_session::VortexSession::with) or rely on -//! [`ArrayKernelsExt::kernels`] to insert the default value. +//! [`ArraySession`](crate::session::ArraySession) owns vortex-array's built-in kernel registry, +//! so sessions that install the default array encodings get their matching built-in kernels too. +//! Sessions can still install a standalone [`ArrayKernels`] registry when they need a kernel-only +//! setup or an explicit override. use std::any::Any; use std::borrow::Borrow; @@ -30,6 +31,7 @@ use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexResult; +use vortex_error::vortex_panic; use vortex_session::SessionExt; use vortex_session::SessionVar; use vortex_session::registry::Id; @@ -46,6 +48,7 @@ use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; +use crate::session::ArraySession; /// Shared hasher used to combine `(outer, child)` tuples into registry keys. static FN_HASHER: LazyLock = LazyLock::new(DefaultHashBuilder::default); @@ -303,9 +306,20 @@ impl SessionVar for ArrayKernels { /// Extension trait for accessing optimizer kernels from a /// [`VortexSession`](vortex_session::VortexSession). pub trait ArrayKernelsExt: SessionExt { - /// Returns the [`ArrayKernels`] session variable. + /// Returns the [`ArrayKernels`] session variable if one is available. + /// + /// A standalone [`ArrayKernels`] variable takes precedence. Otherwise, sessions that include + /// [`ArraySession`] use its built-in kernel registry. + fn kernels_opt(&self) -> Option<&ArrayKernels> { + self.get_opt::() + .or_else(|| self.get_opt::().map(ArraySession::kernels)) + } + + /// Returns the available [`ArrayKernels`] registry. fn kernels(&self) -> &ArrayKernels { - self.get::() + self.kernels_opt().unwrap_or_else(|| { + vortex_panic!("Session contains neither ArrayKernels nor ArraySession") + }) } } diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index ca3ab1218ce..efbafbf85f5 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -21,7 +21,7 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use crate::ArrayRef; -use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ArrayKernelsExt; pub mod kernels; pub mod rules; @@ -68,7 +68,7 @@ fn try_optimize( ) -> VortexResult> { let mut current_array = array.clone(); let mut any_optimizations = false; - let array_ref = session.and_then(|s| s.get_opt::()); + let array_ref = session.and_then(|s| s.kernels_opt()); // Apply reduction rules to the current array until no more rules apply. let mut loop_counter = 0; diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index b63fa64682e..eb6c3f6cfba 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -8,6 +8,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::SessionExt; use vortex_session::SessionVar; +use vortex_session::VortexSession; use vortex_session::registry::Registry; use crate::ArrayRef; @@ -29,6 +30,7 @@ use crate::arrays::Struct; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; +use crate::optimizer::kernels::ArrayKernels; pub type ArrayRegistry = Registry; @@ -36,12 +38,15 @@ pub type ArrayRegistry = Registry; pub struct ArraySession { /// The set of registered array encodings. registry: ArrayRegistry, + /// The set of built-in array kernels for the registered encodings. + kernels: ArrayKernels, } impl ArraySession { pub fn empty() -> ArraySession { Self { registry: ArrayRegistry::default(), + kernels: ArrayKernels::empty(), } } @@ -49,6 +54,10 @@ impl ArraySession { &self.registry } + pub fn kernels(&self) -> &ArrayKernels { + &self.kernels + } + /// Register a new array encoding, replacing any existing encoding with the same ID. pub fn register(&self, plugin: P) { self.registry @@ -60,6 +69,7 @@ impl Default for ArraySession { fn default() -> Self { let this = ArraySession { registry: ArrayRegistry::default(), + kernels: ArrayKernels::default(), }; // Register the canonical encodings. @@ -82,6 +92,9 @@ impl Default for ArraySession { this.register(Masked); this.register(VarBin); + let session = VortexSession::empty().with_some(this.kernels.clone()); + crate::arrays::initialize(&session); + this } } @@ -117,3 +130,29 @@ pub trait ArraySessionExt: SessionExt { } impl ArraySessionExt for S {} + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use crate::ArrayVTable; + use crate::arrays::Bool; + use crate::optimizer::kernels::ArrayKernels; + use crate::optimizer::kernels::ArrayKernelsExt; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::binary::Binary; + use crate::session::ArraySession; + + #[test] + fn array_session_default_registers_builtin_kernels() { + let session = VortexSession::empty().with::(); + + assert!(session.get_opt::().is_none()); + assert!( + session + .kernels() + .find_execute_parent(Binary.id(), Bool.id()) + .is_some() + ); + } +} diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 7f4e72630d8..8b4afceb1fb 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -14,7 +14,6 @@ use vortex_array::dtype::session::DTypeSession; // pulled back out into a vortex_expr crate. pub use vortex_array::expr; use vortex_array::memory::MemorySession; -use vortex_array::optimizer::kernels::ArrayKernels; pub use vortex_array::scalar_fn; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; @@ -171,12 +170,10 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::() - .with::() .with::() .with::() .with::() .with::(); - vortex_array::initialize(&session); #[cfg(feature = "files")] let session = { From 7b7b174e8942b252f50ae39688a5b8fa90d2fadc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 13:50:52 -0400 Subject: [PATCH 12/19] Fix array kernel rustdoc links Signed-off-by: Nicholas Gates --- vortex-array/src/lib.rs | 2 +- vortex-array/src/optimizer/kernels.rs | 2 +- vortex-array/src/optimizer/mod.rs | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 2c0ce752329..fa9fbd9cebb 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -86,7 +86,7 @@ pub mod flatbuffers { } /// Register vortex-array's built-in session-scoped kernels into a standalone -/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. +/// [`ArrayKernels`] registry. /// /// Sessions that use [`ArraySession`] already receive vortex-array's built-in kernels from /// [`ArraySession::default`]. diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index c13544c2b42..7123c1044a5 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -18,7 +18,7 @@ //! Because registered functions have different signatures for each kernel kind, the registry //! maintains one storage map per function type rather than a single type-erased map. //! -//! [`ArraySession`](crate::session::ArraySession) owns vortex-array's built-in kernel registry, +//! [`ArraySession`] owns vortex-array's built-in kernel registry, //! so sessions that install the default array encodings get their matching built-in kernels too. //! Sessions can still install a standalone [`ArrayKernels`] registry when they need a kernel-only //! setup or an explicit override. diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index efbafbf85f5..4170d7b3f24 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -10,7 +10,7 @@ //! There are three public entry points on [`ArrayOptimizer`]: //! //! - [`ArrayOptimizer::optimize`] uses only static rules registered on encoding vtables. -//! - [`ArrayOptimizer::optimize_ctx`] also consults session-scoped [`ArrayKernels`] before +//! - [`ArrayOptimizer::optimize_ctx`] also consults session-scoped [`kernels::ArrayKernels`] before //! static parent-reduce rules, so this is the entry point used by execution. //! - [`ArrayOptimizer::optimize_recursive`] applies the session-aware optimizer to the root and //! every descendant. @@ -31,13 +31,13 @@ pub trait ArrayOptimizer { /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint. /// /// This uses only static rules registered on encoding vtables. Use [`Self::optimize_ctx`] - /// when session-registered [`ArrayKernels`] should participate. + /// when session-registered [`kernels::ArrayKernels`] should participate. fn optimize(&self) -> VortexResult; - /// Optimize the root array node using static rules and any [`ArrayKernels`] on `session`. + /// Optimize the root array node using static rules and any [`kernels::ArrayKernels`] on `session`. /// /// Session kernels are checked for each `(parent_encoding_id, child_encoding_id)` pair before - /// the child's static `PARENT_RULES`. If `session` does not contain [`ArrayKernels`], this + /// the child's static `PARENT_RULES`. If `session` does not contain [`kernels::ArrayKernels`], this /// behaves like [`Self::optimize`]. fn optimize_ctx(&self, session: &VortexSession) -> VortexResult; From d8ae50e40f5e90534b41121ec352b46269def10f Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 14:59:33 -0400 Subject: [PATCH 13/19] Address execute parent review findings Signed-off-by: Nicholas Gates --- docs/developer-guide/internals/execution.md | 13 ++++++++++--- vortex-array/src/executor.rs | 10 ++++++++++ vortex-array/src/optimizer/kernels.rs | 3 +++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/docs/developer-guide/internals/execution.md b/docs/developer-guide/internals/execution.md index be92e29ceef..de8aa48ff91 100644 --- a/docs/developer-guide/internals/execution.md +++ b/docs/developer-guide/internals/execution.md @@ -115,8 +115,9 @@ Each child is given the opportunity to execute its parent in a fused manner via `ExecuteParentKernel`. Unlike reduce rules, parent kernels may read buffers and perform real computation. -An encoding declares its parent kernels in a `ParentKernelSet`, specifying which parent types -each kernel handles via a `Matcher`: +An encoding declares parent kernels by implementing `ExecuteParentKernel` and registering them +with the session-scoped kernel registry, specifying which parent types each kernel handles via a +`Matcher`: ```rust pub trait ExecuteParentKernel { @@ -124,12 +125,18 @@ pub trait ExecuteParentKernel { fn execute_parent( &self, - array: &V::Array, // the child + array: ArrayView<'_, V>, // the child parent: ::Match<'_>, // the matched parent child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult>; } + +pub fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(parent_id, Child, Kernel); +} ``` Examples: diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 600d07d6d52..ded6b8a6aa6 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -569,6 +569,16 @@ fn execute_parent_for_child( if let Some(plugins) = kernels.get(&key) { for plugin in plugins.as_ref() { if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { + if cfg!(debug_assertions) { + vortex_ensure!( + result.len() == parent.len(), + "Executed parent canonical length mismatch" + ); + vortex_ensure!( + result.dtype() == parent.dtype(), + "Executed parent canonical dtype mismatch" + ); + } return Ok(Some(result)); } } diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 7123c1044a5..3ddb8d08f7d 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -251,6 +251,9 @@ impl ArrayKernels { /// The executor invokes registered kernels in registration order before falling through to /// later registered kernels for the same key. `parent` is usually the parent array's encoding /// id. For `ScalarFnArray`, it is the scalar function id, for example `Cast.id()`. + /// + /// If kernels have already been registered for the same pair, this kernel is appended after + /// them; registering for an existing key cannot override built-in kernels installed earlier. pub fn register_execute_parent_kernel(&self, parent: Id, child: V, kernel: K) where V: VTable, From d125dd059e51166cbf3d142e088e51b3f47ca7d9 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 15:17:50 -0400 Subject: [PATCH 14/19] Update execute parent docs diagram Signed-off-by: Nicholas Gates --- docs/developer-guide/internals/execution.md | 29 ++++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/docs/developer-guide/internals/execution.md b/docs/developer-guide/internals/execution.md index de8aa48ff91..15ba16a8135 100644 --- a/docs/developer-guide/internals/execution.md +++ b/docs/developer-guide/internals/execution.md @@ -202,17 +202,23 @@ execute_until(root): │ restore builder, │ │ loop ▼ │ ┌────────────────────────────────────────────┐ - │ │ Step 2a: current_array.execute_parent( │ - │ │ stack.top.parent_array ) │ - │ │ child looks UP at the suspended parent │ + │ │ Step 2a: execute_parent_for_child( │ + │ │ stack.top.parent_array, │ + │ │ current_array ) │ + │ │ lookup session kernels by key: │ + │ │ (parent.encoding_id(), │ + │ │ child.encoding_id()) │ │ ├────────────┬───────────────────────────────┘ │ │ Some │ None │ │ │ │ │ ▼ │ │ ┌─────────────────────────────────────────┐ - │ │ │ Step 2b: each child.execute_parent( │ - │ │ │ current_array ) │ - │ │ │ children look UP at current_array │ + │ │ │ Step 2b: for each child: │ + │ │ │ execute_parent_for_child( │ + │ │ │ current_array, child ) │ + │ │ │ lookup session kernels by key: │ + │ │ │ (parent.encoding_id(), │ + │ │ │ child.encoding_id()) │ │ │ ├──────────┬──────────────────────────────┘ │ │ │ Some │ None │ │ │ │ @@ -243,6 +249,9 @@ Step 2a and Step 2b are skipped while `current_builder` is active. `AppendChild` consumes `current_array`: some slots already live in the builder, so a parent rewrite would observe inconsistent state and could discard accumulated builder data. +Both Step 2a and Step 2b route through `execute_parent_for_child`, which looks up +`(parent.encoding_id(), child.encoding_id())` in the session kernel snapshot. + ## Incremental Execution Execution is incremental: each call to `execute` moves the array one step closer to canonical @@ -263,12 +272,12 @@ codes through the slice, missing the Dict-RLE optimization entirely. Incremental avoids this: 1. First iteration: the slice `execute` returns `ExecuteSlot` for its `RunEndArray` child. - Once that child is in focus, Step 2a gives it a chance to rewrite the suspended slice - parent before the child is forced toward canonical form. + Once that child is in focus, Step 2a looks up a registered parent kernel for the + `(slice, runend)` pair before the child is forced toward canonical form. 2. Second iteration: the `RunEndArray` codes child now matches the Dict-RLE pattern. Its - `execute_parent` provides a fused kernel that expands runs while performing dictionary - lookups in a single pass, returning the canonical array directly. + registered execute-parent kernel expands runs while performing dictionary lookups in a single + pass, returning the canonical array directly. ## Walkthrough: Executing a RunEnd-Encoded Array From b7a1dd9a03341df544b632d9cb233cf13cbf2992 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 15:36:19 -0400 Subject: [PATCH 15/19] Document execute parent registry semantics Signed-off-by: Nicholas Gates --- vortex-array/src/executor.rs | 54 +++++++++++++++++++++++++++ vortex-array/src/lib.rs | 14 ++++--- vortex-array/src/optimizer/kernels.rs | 24 +++++------- vortex-array/src/optimizer/mod.rs | 15 +++++--- vortex-array/src/session/mod.rs | 29 +++++++++++++- vortex-file/src/lib.rs | 25 +++++++++++++ 6 files changed, 133 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index ded6b8a6aa6..50173f95ff9 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -328,6 +328,10 @@ pub struct ExecutionCtx { impl ExecutionCtx { /// Create a new execution context with the given session. + /// + /// This captures a snapshot of the session's execute-parent kernel registry. Kernels + /// registered after this context is created are not visible to it; create a new + /// [`ExecutionCtx`] after registration to use newly registered kernels. pub fn new(session: VortexSession) -> Self { let execute_parent_kernels = session .kernels_opt() @@ -847,3 +851,53 @@ impl VortexSessionExecute for VortexSession { ExecutionCtx::new(self.clone()) } } + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use super::*; + use crate::VTable as _; + use crate::arrays::Bool; + use crate::arrays::Primitive; + use crate::optimizer::kernels::ArrayKernels; + use crate::optimizer::kernels::ExecuteParentFn; + use crate::optimizer::kernels::execute_parent_key; + + fn noop_execute_parent( + _child: &ArrayRef, + _parent: &ArrayRef, + _child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } + + #[test] + fn execution_ctx_snapshots_execute_parent_kernels_at_creation() { + let session = VortexSession::empty().with::(); + let key = execute_parent_key(Bool.id(), Primitive.id()); + + let before_registration = ExecutionCtx::new(session.clone()); + assert!( + !before_registration + .execute_parent_kernels + .contains_key(&key) + ); + + session.kernels().register_execute_parent( + Bool.id(), + Primitive.id(), + &[noop_execute_parent as ExecuteParentFn], + ); + + assert!( + !before_registration + .execute_parent_kernels + .contains_key(&key) + ); + + let after_registration = ExecutionCtx::new(session); + assert!(after_registration.execute_parent_kernels.contains_key(&key)); + } +} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index fa9fbd9cebb..151e166b1ba 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -31,7 +31,7 @@ use crate::aggregate_fn::session::AggregateFnSession; use crate::arrow::ArrowSession; use crate::dtype::session::DTypeSession; use crate::memory::MemorySession; -use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::session::ScalarFnSession; use crate::session::ArraySession; use crate::stats::session::StatsSession; @@ -85,13 +85,15 @@ pub mod flatbuffers { pub use vortex_flatbuffers::array::*; } -/// Register vortex-array's built-in session-scoped kernels into a standalone -/// [`ArrayKernels`] registry. +/// Register vortex-array's built-in session-scoped kernels into the active +/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. /// -/// Sessions that use [`ArraySession`] already receive vortex-array's built-in kernels from -/// [`ArraySession::default`]. +/// If the session contains a standalone [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels), +/// this registers into that registry. Otherwise, if the session contains an [`ArraySession`], this +/// registers into the [`ArraySession`]-owned registry. Sessions that use [`ArraySession::default`] +/// already receive these built-in kernels. pub fn initialize(session: &VortexSession) { - if session.get_opt::().is_some() { + if session.kernels_opt().is_some() { arrays::initialize(session); } } diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 3ddb8d08f7d..9363488df65 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -269,16 +269,11 @@ impl ArrayKernels { ); } - /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`. - /// - /// Returns an owned [`Arc`] so the session-variable borrow can be dropped before invoking the - /// functions. - pub fn find_execute_parent( - &self, - parent: Id, - child: Id, - ) -> Option> { - self.execute_parent.get(&hash_fn_id(parent, child)) + /// Returns true when one or more execute-parent kernels are registered for `(parent, child)`. + pub fn has_execute_parent(&self, parent: Id, child: Id) -> bool { + self.execute_parent + .get(&hash_fn_id(parent, child)) + .is_some() } /// Return the currently published execute-parent kernel snapshot. @@ -309,16 +304,17 @@ impl SessionVar for ArrayKernels { /// Extension trait for accessing optimizer kernels from a /// [`VortexSession`](vortex_session::VortexSession). pub trait ArrayKernelsExt: SessionExt { - /// Returns the [`ArrayKernels`] session variable if one is available. + /// Returns the active [`ArrayKernels`] registry if one is available. /// - /// A standalone [`ArrayKernels`] variable takes precedence. Otherwise, sessions that include - /// [`ArraySession`] use its built-in kernel registry. + /// A standalone [`ArrayKernels`] variable takes precedence and is not merged with an + /// [`ArraySession`]-owned registry. Otherwise, sessions that include [`ArraySession`] use its + /// built-in kernel registry. fn kernels_opt(&self) -> Option<&ArrayKernels> { self.get_opt::() .or_else(|| self.get_opt::().map(ArraySession::kernels)) } - /// Returns the available [`ArrayKernels`] registry. + /// Returns the active [`ArrayKernels`] registry. fn kernels(&self) -> &ArrayKernels { self.kernels_opt().unwrap_or_else(|| { vortex_panic!("Session contains neither ArrayKernels nor ArraySession") diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 4170d7b3f24..4bbec8b7c32 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -10,8 +10,9 @@ //! There are three public entry points on [`ArrayOptimizer`]: //! //! - [`ArrayOptimizer::optimize`] uses only static rules registered on encoding vtables. -//! - [`ArrayOptimizer::optimize_ctx`] also consults session-scoped [`kernels::ArrayKernels`] before -//! static parent-reduce rules, so this is the entry point used by execution. +//! - [`ArrayOptimizer::optimize_ctx`] also consults the session's active +//! [`kernels::ArrayKernels`] registry before static parent-reduce rules, so this is the entry +//! point used by execution. //! - [`ArrayOptimizer::optimize_recursive`] applies the session-aware optimizer to the root and //! every descendant. @@ -31,14 +32,16 @@ pub trait ArrayOptimizer { /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint. /// /// This uses only static rules registered on encoding vtables. Use [`Self::optimize_ctx`] - /// when session-registered [`kernels::ArrayKernels`] should participate. + /// when a session-scoped [`kernels::ArrayKernels`] registry should participate. fn optimize(&self) -> VortexResult; - /// Optimize the root array node using static rules and any [`kernels::ArrayKernels`] on `session`. + /// Optimize the root array node using static rules and the active + /// [`kernels::ArrayKernels`] registry on `session`, if any. /// /// Session kernels are checked for each `(parent_encoding_id, child_encoding_id)` pair before - /// the child's static `PARENT_RULES`. If `session` does not contain [`kernels::ArrayKernels`], this - /// behaves like [`Self::optimize`]. + /// the child's static `PARENT_RULES`. A standalone [`kernels::ArrayKernels`] registry takes + /// precedence over an [`ArraySession`](crate::session::ArraySession)-owned registry. If + /// `session` contains neither registry, this behaves like [`Self::optimize`]. fn optimize_ctx(&self, session: &VortexSession) -> VortexResult; /// Optimize the entire array tree recursively (root and all descendants). diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index eb6c3f6cfba..eec38d4c15d 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -54,6 +54,10 @@ impl ArraySession { &self.registry } + /// Returns this array session's built-in kernel registry. + /// + /// This registry is used when the surrounding [`VortexSession`] does not contain a standalone + /// [`ArrayKernels`] session variable. A standalone [`ArrayKernels`] shadows this registry. pub fn kernels(&self) -> &ArrayKernels { &self.kernels } @@ -148,11 +152,32 @@ mod tests { let session = VortexSession::empty().with::(); assert!(session.get_opt::().is_none()); + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn initialize_registers_builtin_kernels_into_empty_array_session() { + let session = VortexSession::empty().with_some(ArraySession::empty()); + + assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); + + crate::initialize(&session); + + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn standalone_array_kernels_shadow_array_session_kernels() { + let session = VortexSession::empty() + .with::() + .with::(); + assert!( session + .get::() .kernels() - .find_execute_parent(Binary.id(), Bool.id()) - .is_some() + .has_execute_parent(Binary.id(), Bool.id()) ); + assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); } } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index dcbf7a6ddb3..9daa55efd79 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -181,3 +181,28 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(feature = "unstable_encodings")] vortex_tensor::initialize(session); } + +#[cfg(test)] +mod default_encoding_tests { + use vortex_array::VTable as _; + use vortex_array::array_session; + use vortex_array::arrays::Filter; + use vortex_array::optimizer::kernels::ArrayKernelsExt as _; + use vortex_array::session::ArraySessionExt as _; + use vortex_fsst::FSST; + + use crate::register_default_encodings; + + #[test] + fn register_default_encodings_registers_external_execute_parent_kernels() { + let session = array_session(); + + assert!(session.arrays().registry().find(&FSST.id()).is_none()); + assert!(!session.kernels().has_execute_parent(Filter.id(), FSST.id())); + + register_default_encodings(&session); + + assert!(session.arrays().registry().find(&FSST.id()).is_some()); + assert!(session.kernels().has_execute_parent(Filter.id(), FSST.id())); + } +} From 5ef7a33409c12174c44dd0ed4322c9c9aa8acf6b Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 19 Jun 2026 14:57:22 -0400 Subject: [PATCH 16/19] Make scalar function validity explicit Signed-off-by: Nicholas Gates --- .../src/arrays/scalar_fn/vtable/mod.rs | 39 ++++++- .../src/arrays/scalar_fn/vtable/validity.rs | 101 ++++++++---------- vortex-array/src/scalar_fn/erased.rs | 8 +- vortex-array/src/scalar_fn/fns/between/mod.rs | 4 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 85 +++++++++++++-- vortex-array/src/scalar_fn/fns/byte_length.rs | 8 +- vortex-array/src/scalar_fn/fns/case_when.rs | 25 +++++ vortex-array/src/scalar_fn/fns/cast/mod.rs | 6 +- vortex-array/src/scalar_fn/fns/dynamic.rs | 9 ++ .../src/scalar_fn/fns/fill_null/mod.rs | 4 +- vortex-array/src/scalar_fn/fns/get_item.rs | 9 ++ vortex-array/src/scalar_fn/fns/is_not_null.rs | 20 +++- vortex-array/src/scalar_fn/fns/is_null.rs | 19 +++- vortex-array/src/scalar_fn/fns/like/mod.rs | 4 +- .../src/scalar_fn/fns/list_contains/mod.rs | 12 +++ vortex-array/src/scalar_fn/fns/literal.rs | 8 +- vortex-array/src/scalar_fn/fns/mask/mod.rs | 6 +- vortex-array/src/scalar_fn/fns/merge.rs | 4 +- vortex-array/src/scalar_fn/fns/not/mod.rs | 9 ++ vortex-array/src/scalar_fn/fns/pack.rs | 4 +- vortex-array/src/scalar_fn/fns/root.rs | 9 ++ vortex-array/src/scalar_fn/fns/select.rs | 8 ++ vortex-array/src/scalar_fn/fns/stat.rs | 9 ++ .../src/scalar_fn/fns/variant_get/mod.rs | 9 ++ vortex-array/src/scalar_fn/fns/zip/mod.rs | 13 +++ vortex-array/src/scalar_fn/foreign.rs | 11 ++ .../src/scalar_fn/internal/row_count.rs | 9 ++ vortex-array/src/scalar_fn/typed.rs | 4 +- vortex-array/src/scalar_fn/vtable.rs | 10 +- 29 files changed, 347 insertions(+), 119 deletions(-) diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index 6833848deae..e350e5cd0c1 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -8,6 +8,7 @@ use std::hash::Hash; use std::hash::Hasher; use std::marker::PhantomData; use std::ops::Deref; +use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexResult; @@ -37,13 +38,18 @@ use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::expr::Expression; +use crate::expr::lit; use crate::matcher::Matcher; use crate::scalar_fn; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ReduceCtx; +use crate::scalar_fn::ReduceNode; +use crate::scalar_fn::ReduceNodeRef; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::TypedScalarFnInstance; use crate::scalar_fn::VecExecutionArgs; use crate::serde::ArrayChildren; @@ -175,7 +181,7 @@ pub trait ScalarFnFactoryExt: scalar_fn::ScalarFnVTable { options: Self::Options, children: impl Into>, ) -> VortexResult { - let scalar_fn = scalar_fn::TypedScalarFnInstance::new(self.clone(), options).erased(); + let scalar_fn = TypedScalarFnInstance::new(self.clone(), options).erased(); let children = children.into(); vortex_ensure!( @@ -201,6 +207,24 @@ pub trait ScalarFnFactoryExt: scalar_fn::ScalarFnVTable { } impl ScalarFnFactoryExt for V {} +pub(crate) fn scalar_fn_array_expr(array: ArrayView<'_, ScalarFn>) -> VortexResult { + let inputs: Vec<_> = array + .iter_children() + .map(|child| { + if let Some(scalar) = child.as_constant() { + return Ok(lit(scalar)); + } + + Expression::try_new( + TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(), + [], + ) + }) + .collect::>()?; + + Expression::try_new(array.scalar_fn().clone(), inputs) +} + /// A matcher that matches any scalar function expression. #[derive(Debug)] pub struct AnyScalarFn; @@ -320,12 +344,21 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr { crate::Executable::execute(options.0.clone(), ctx) } + fn reduce( + &self, + options: &Self::Options, + _node: &dyn ReduceNode, + _ctx: &dyn ReduceCtx, + ) -> VortexResult> { + Ok(Some(Arc::new(options.0.clone()))) + } + fn validity( &self, options: &Self::Options, _expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { let validity_array = options.0.validity()?.to_array(options.0.len()); - Ok(Some(ArrayExpr.new_expr(FakeEq(validity_array), []))) + Ok(ArrayExpr.new_expr(FakeEq(validity_array), [])) } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 2ac376155e3..5a738b1a688 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -3,72 +3,57 @@ use vortex_error::VortexResult; -use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::array::ValidityVTable; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::vtable::ArrayExpr; -use crate::arrays::scalar_fn::vtable::FakeEq; +use crate::arrays::ConstantArray; use crate::arrays::scalar_fn::vtable::ScalarFn; -use crate::expr::Expression; -use crate::expr::lit; -use crate::scalar_fn::TypedScalarFnInstance; -use crate::scalar_fn::VecExecutionArgs; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::root::Root; +use crate::arrays::scalar_fn::vtable::scalar_fn_array_expr; use crate::validity::Validity; -/// Execute an expression tree recursively. -/// -/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals. -fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - - // Handle Root expression - this should not happen in validity expressions - if expr.is::() { - vortex_error::vortex_bail!("Root expression cannot be executed in validity context"); - } - - // Handle Literal expression - create a constant array - if expr.is::() { - let scalar = expr.as_::(); - return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array()); - } - - // Recursively execute child expressions to get input arrays - let inputs: Vec = expr - .children() - .iter() - .map(|child| execute_expr(child, row_count)) - .collect::>()?; - - let args = VecExecutionArgs::new(inputs, row_count); - - Ok(expr.scalar_fn().execute(&args, &mut ctx)?.into_array()) -} - impl ValidityVTable for ScalarFn { fn validity(array: ArrayView<'_, ScalarFn>) -> VortexResult { - let inputs: Vec<_> = array - .iter_children() - .map(|child| { - if let Some(scalar) = child.as_constant() { - return Ok(lit(scalar)); - } - Expression::try_new( - TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(), - [], - ) - }) - .collect::>()?; - - let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?; - let validity_expr = array.scalar_fn().validity(&expr)?; + let expr = scalar_fn_array_expr(array)?.validity()?; + let input = ConstantArray::new(true, array.len()).into_array(); + Ok(Validity::Array(input.apply(&expr)?)) + } +} - // Execute the validity expression. All leaves are ArrayExpr nodes. - Ok(Validity::Array(execute_expr(&validity_expr, array.len())?)) +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use crate::IntoArray; + use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + use crate::arrays::ScalarFn; + use crate::arrays::scalar_fn::ScalarFnArrayExt; + use crate::arrays::scalar_fn::vtable::ScalarFnFactoryExt; + use crate::scalar_fn::fns::binary::Binary; + use crate::scalar_fn::fns::operators::Operator; + use crate::validity::Validity; + + #[test] + fn scalar_fn_validity_stays_lazy() -> VortexResult<()> { + let lhs = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); + let rhs = BoolArray::from_iter([Some(true), Some(false), None]).into_array(); + let predicate = Binary.try_new_array(lhs.len(), Operator::And, [lhs, rhs])?; + + let Validity::Array(validity_array) = predicate.validity()? else { + panic!("scalar function validity should be represented as an array"); + }; + + let validity_scalar_fn = validity_array + .as_opt::() + .expect("validity should remain a lazy scalar function array"); + assert!(validity_scalar_fn.scalar_fn().is::()); + + let validity_mask = + validity_array.execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + assert!(validity_mask.all_true()); + + Ok(()) } } diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 10e82d25455..2a33114ccd5 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -22,16 +22,13 @@ use crate::dtype::DType; use crate::expr::Expression; use crate::expr::StatsCatalog; use crate::expr::stats::Stat; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ReduceCtx; use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ReduceNodeRef; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::SimplifyCtx; -use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::options::ScalarFnOptions; use crate::scalar_fn::signature::ScalarFnSignature; use crate::scalar_fn::typed::DynScalarFn; @@ -134,10 +131,7 @@ impl ScalarFnRef { /// Transforms the expression into one representing the validity of this expression. pub fn validity(&self, expr: &Expression) -> VortexResult { - Ok(self.0.validity(expr)?.unwrap_or_else(|| { - // TODO(ngates): make validity a mandatory method on VTable to avoid this fallback. - IsNotNull.new_expr(EmptyOptions, [expr.clone()]) - })) + self.0.validity(expr) } /// Execute the expression given the input arguments. diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 013438e23b2..430534c41cf 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -318,11 +318,11 @@ impl ScalarFnVTable for Between { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { let arr = expression.child(0).validity()?; let lower = expression.child(1).validity()?; let upper = expression.child(2).validity()?; - Ok(Some(and(and(arr, lower), upper))) + Ok(and(and(arr, lower), upper)) } fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index bbb392a3752..c737131333a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -22,11 +22,14 @@ use crate::expr::and; use crate::expr::and_collect; use crate::expr::eq; use crate::expr::expression::Expression; +use crate::expr::fill_null; use crate::expr::gt; use crate::expr::gt_eq; use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; +use crate::expr::not; +use crate::expr::or; use crate::expr::or_collect; use crate::expr::stats::Stat; use crate::scalar_fn::Arity; @@ -273,21 +276,30 @@ impl ScalarFnVTable for Binary { } } - fn validity( - &self, - operator: &Operator, - expression: &Expression, - ) -> VortexResult> { + fn validity(&self, operator: &Operator, expression: &Expression) -> VortexResult { let lhs = expression.child(0).validity()?; let rhs = expression.child(1).validity()?; Ok(match operator { - // AND and OR are kleene logic. - Operator::And => None, - Operator::Or => None, + // AND and OR are Kleene logic. Their result is valid if both children are valid, + // or if a valid child value alone determines the result. + Operator::And => or( + and(lhs, rhs), + or( + not(fill_null(expression.child(0).clone(), lit(true))), + not(fill_null(expression.child(1).clone(), lit(true))), + ), + ), + Operator::Or => or( + and(lhs, rhs), + or( + fill_null(expression.child(0).clone(), lit(false)), + fill_null(expression.child(1).clone(), lit(false)), + ), + ), _ => { // All other binary operators are null if either side is null. - Some(and(lhs, rhs)) + and(lhs, rhs) } }) } @@ -318,6 +330,7 @@ impl ScalarFnVTable for Binary { #[cfg(test)] mod tests { use vortex_error::VortexExpect; + use vortex_error::VortexResult; use super::*; use crate::LEGACY_SESSION; @@ -544,6 +557,60 @@ mod tests { assert_arrays_eq!(result, BoolArray::from_iter([Some(true)]).into_array()) } + #[test] + fn test_kleene_boolean_result_validity() -> VortexResult<()> { + use crate::IntoArray; + use crate::arrays::BoolArray; + use crate::validity::Validity; + + let lhs = BoolArray::from_iter([ + Some(true), + Some(true), + Some(false), + Some(false), + None, + None, + Some(true), + Some(false), + None, + ]) + .into_array(); + let rhs = BoolArray::from_iter([ + Some(true), + Some(false), + Some(true), + None, + Some(true), + Some(false), + None, + None, + None, + ]) + .into_array(); + + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + let and_result = lhs.binary(rhs.clone(), Operator::And)?; + let and_expected = + Validity::from_iter([true, true, true, true, false, true, false, true, false]); + assert!( + and_result + .validity()? + .mask_eq(&and_expected, and_result.len(), &mut ctx)? + ); + + let or_result = lhs.binary(rhs, Operator::Or)?; + let or_expected = + Validity::from_iter([true, true, true, false, true, false, true, false, false]); + assert!( + or_result + .validity()? + .mask_eq(&or_expected, or_result.len(), &mut ctx)? + ); + + Ok(()) + } + #[test] fn test_scalar_subtract_unsigned() { use vortex_buffer::buffer; diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index aa9c508ea89..5b2de882a9d 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -123,12 +123,8 @@ impl ScalarFnVTable for ByteLength { } } - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - Ok(Some(expression.child(0).validity()?)) + fn validity(&self, _: &Self::Options, expression: &Expression) -> VortexResult { + expression.child(0).validity() } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 0a5f6b4ee14..92c33a4c7b8 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -35,6 +35,8 @@ use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::lit; +use crate::expr::zip_expr; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -300,6 +302,29 @@ impl ScalarFnVTable for CaseWhen { Ok(Some(crate::expr::fill_null(x.clone(), fill.clone()))) } + fn validity( + &self, + options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + let num_pairs = options.num_when_then_pairs as usize; + let mut validity = if options.has_else { + expression.child(num_pairs * 2).validity()? + } else { + lit(false) + }; + + for pair_idx in (0..num_pairs).rev() { + validity = zip_expr( + expression.child(pair_idx * 2).clone(), + expression.child(pair_idx * 2 + 1).validity()?, + validity, + ); + } + + Ok(validity) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { true } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index abc59af2c9a..364c2da80a5 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -185,12 +185,12 @@ impl ScalarFnVTable for Cast { } } - fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult> { - Ok(Some(if dtype.is_nullable() { + fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult { + Ok(if dtype.is_nullable() { expression.child(0).validity()? } else { lit(true) - })) + }) } // This might apply a nullability diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 7efebf79220..0a0ef0834ac 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -21,6 +21,7 @@ use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::Expression; use crate::expr::StatsCatalog; +use crate::expr::is_not_null; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; @@ -164,6 +165,14 @@ impl ScalarFnVTable for DynamicComparison { } } + fn validity( + &self, + _dynamic: &DynamicComparisonExpr, + expression: &Expression, + ) -> VortexResult { + Ok(is_not_null(expression.clone())) + } + // Defer to the child fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { false diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index b2d07b41255..6b19fef6f51 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -126,10 +126,10 @@ impl ScalarFnVTable for FillNull { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { // After fill_null, the result validity depends on the fill value's nullability. // If fill_value is non-nullable, the result is always valid. - Ok(Some(expression.child(1).validity()?)) + expression.child(1).validity() } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index de7e45ca9b0..cfc77fc8e0b 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -22,6 +22,7 @@ use crate::dtype::FieldPath; use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::StatsCatalog; +use crate::expr::is_not_null; use crate::expr::lit; use crate::expr::stats::Stat; use crate::scalar_fn::Arity; @@ -188,6 +189,14 @@ impl ScalarFnVTable for GetItem { Ok(None) } + fn validity( + &self, + _field_name: &FieldName, + expression: &Expression, + ) -> VortexResult { + Ok(is_not_null(expression.clone())) + } + fn stat_expression( &self, field_name: &FieldName, diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 589333304e2..bdd42fa84ff 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -12,11 +12,13 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::arrays::ScalarFn; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::StatsCatalog; use crate::expr::eq; +use crate::expr::lit; use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -81,10 +83,16 @@ impl ScalarFnVTable for IsNotNull { &self, _data: &Self::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let child = args.get(0)?; - match child.validity()? { + let validity = if child.is::() { + child.execute::(ctx)?.validity()? + } else { + child.validity()? + }; + + match validity { Validity::NonNullable | Validity::AllValid => { Ok(ConstantArray::new(true, args.row_count()).into_array()) } @@ -93,6 +101,14 @@ impl ScalarFnVTable for IsNotNull { } } + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } + fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { true } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 7315fbe8c07..0d9edb84564 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -9,6 +9,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::arrays::ScalarFn; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -68,14 +69,20 @@ impl ScalarFnVTable for IsNull { &self, _data: &Self::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let child = args.get(0)?; if let Some(scalar) = child.as_constant() { return Ok(ConstantArray::new(scalar.is_null(), args.row_count()).into_array()); } - match child.validity()? { + let validity = if child.is::() { + child.execute::(ctx)?.validity()? + } else { + child.validity()? + }; + + match validity { Validity::NonNullable | Validity::AllValid => { Ok(ConstantArray::new(false, args.row_count()).into_array()) } @@ -84,6 +91,14 @@ impl ScalarFnVTable for IsNull { } } + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } + fn stat_falsification( &self, _options: &Self::Options, diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index b7f357020f1..899f216d58a 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -155,11 +155,11 @@ impl ScalarFnVTable for Like { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { tracing::warn!("Computing validity for LIKE expression"); let child_validity = expression.child(0).validity()?; let pattern_validity = expression.child(1).validity()?; - Ok(Some(and(child_validity, pattern_validity))) + Ok(and(child_validity, pattern_validity)) } fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index e16991763ed..56182c166c5 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -35,6 +35,7 @@ use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::StatsCatalog; +use crate::expr::and; use crate::expr::and_collect; use crate::expr::gt; use crate::expr::lit; @@ -166,6 +167,17 @@ impl ScalarFnVTable for ListContains { None } + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + Ok(and( + expression.child(0).validity()?, + expression.child(1).validity()?, + )) + } + // Nullability matters for contains([], x) where x is false. fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { true diff --git a/vortex-array/src/scalar_fn/fns/literal.rs b/vortex-array/src/scalar_fn/fns/literal.rs index 16b112e5a78..9ed04121048 100644 --- a/vortex-array/src/scalar_fn/fns/literal.rs +++ b/vortex-array/src/scalar_fn/fns/literal.rs @@ -140,12 +140,8 @@ impl ScalarFnVTable for Literal { } } - fn validity( - &self, - scalar: &Scalar, - _expression: &Expression, - ) -> VortexResult> { - Ok(Some(lit(scalar.is_valid()))) + fn validity(&self, scalar: &Scalar, _expression: &Expression) -> VortexResult { + Ok(lit(scalar.is_valid())) } fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/fns/mask/mod.rs index 4dd55948ec8..44baab5f6f1 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mask/mod.rs @@ -127,11 +127,11 @@ impl ScalarFnVTable for Mask { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { - Ok(Some(and( + ) -> VortexResult { + Ok(and( expression.child(0).validity()?, expression.child(1).clone(), - ))) + )) } } diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 390d3a2213e..94355c24de1 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -237,8 +237,8 @@ impl ScalarFnVTable for Merge { &self, _options: &Self::Options, _expression: &Expression, - ) -> VortexResult> { - Ok(Some(lit(true))) + ) -> VortexResult { + Ok(lit(true)) } fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/not/mod.rs b/vortex-array/src/scalar_fn/fns/not/mod.rs index 156a6568df7..5a0014c7d57 100644 --- a/vortex-array/src/scalar_fn/fns/not/mod.rs +++ b/vortex-array/src/scalar_fn/fns/not/mod.rs @@ -18,6 +18,7 @@ use crate::arrays::ConstantArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::expr::Expression; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -98,6 +99,14 @@ impl ScalarFnVTable for Not { child.execute::(ctx)?.not() } + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + expression.child(0).validity() + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { false } diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 2a87bf069bc..b2f81f84ac3 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -130,8 +130,8 @@ impl ScalarFnVTable for Pack { &self, _options: &Self::Options, _expression: &Expression, - ) -> VortexResult> { - Ok(Some(lit(true))) + ) -> VortexResult { + Ok(lit(true)) } fn execute( diff --git a/vortex-array/src/scalar_fn/fns/root.rs b/vortex-array/src/scalar_fn/fns/root.rs index 87b8b62ccf4..eeaf02bb832 100644 --- a/vortex-array/src/scalar_fn/fns/root.rs +++ b/vortex-array/src/scalar_fn/fns/root.rs @@ -14,6 +14,7 @@ use crate::dtype::DType; use crate::dtype::FieldPath; use crate::expr::StatsCatalog; use crate::expr::expression::Expression; +use crate::expr::is_not_null; use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -80,6 +81,14 @@ impl ScalarFnVTable for Root { vortex_bail!("Root expression is not executable") } + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + Ok(is_not_null(expression.clone())) + } + fn stat_expression( &self, _options: &Self::Options, diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index 54b63a00f89..c4c22cfd5cd 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -232,6 +232,14 @@ impl ScalarFnVTable for Select { Ok(None) } + fn validity( + &self, + _selection: &FieldSelection, + expression: &Expression, + ) -> VortexResult { + expression.child(0).validity() + } + fn is_null_sensitive(&self, _instance: &FieldSelection) -> bool { true } diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 84fc5760495..d3f39af7efb 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -21,6 +21,7 @@ use crate::aggregate_fn::fns::all_null::AllNull; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::is_not_null; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -123,6 +124,14 @@ impl ScalarFnVTable for StatFn { let dtype = stat_dtype(options.aggregate_fn(), input.dtype())?; stat_array(&input, options.aggregate_fn(), dtype, args.row_count()) } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + Ok(is_not_null(expression.clone())) + } } fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult { diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index ac6aa562aa6..aa292332fdf 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -26,6 +26,7 @@ use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; use crate::expr::Expression; +use crate::expr::is_not_null; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -163,6 +164,14 @@ impl ScalarFnVTable for VariantGet { let array = ChunkedArray::try_new(chunks, dtype)?.into_array(); VariantArray::try_new(array, None).map(|array| array.into_array()) } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + Ok(is_not_null(expression.clone())) + } } fn variant_get_scalar( diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 8a3e1b5bbc6..0f5d4bb96e9 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -24,6 +24,7 @@ use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::zip_expr; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -155,6 +156,18 @@ impl ScalarFnVTable for Zip { Ok(None) } + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult { + Ok(zip_expr( + expression.child(2).clone(), + expression.child(0).validity()?, + expression.child(1).validity()?, + )) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { true } diff --git a/vortex-array/src/scalar_fn/foreign.rs b/vortex-array/src/scalar_fn/foreign.rs index d2abd3b9a95..eefe799a279 100644 --- a/vortex-array/src/scalar_fn/foreign.rs +++ b/vortex-array/src/scalar_fn/foreign.rs @@ -118,4 +118,15 @@ impl ScalarFnVTable for ForeignScalarFnVTable { ) -> VortexResult { vortex_bail!("Cannot execute unknown scalar function '{}'", self.id); } + + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + vortex_bail!( + "Cannot compute validity for unknown scalar function '{}'", + self.id + ); + } } diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs index eee838f803d..6429502a5b6 100644 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -12,6 +12,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::Expression; +use vortex_array::expr::lit; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -83,6 +84,14 @@ impl ScalarFnVTable for RowCount { vortex_bail!("RowCount must be substituted before evaluation") } + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { false } diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index fbfc8caace4..16f64eb2e51 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -102,7 +102,7 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { ctx: &dyn SimplifyCtx, ) -> VortexResult>; fn simplify_untyped(&self, expression: &Expression) -> VortexResult>; - fn validity(&self, expression: &Expression) -> VortexResult>; + fn validity(&self, expression: &Expression) -> VortexResult; fn stat_falsification( &self, expression: &Expression, @@ -221,7 +221,7 @@ impl DynScalarFn for TypedScalarFnInstance { V::simplify_untyped(&self.vtable, &self.options, expression) } - fn validity(&self, expression: &Expression) -> VortexResult> { + fn validity(&self, expression: &Expression) -> VortexResult { V::validity(&self.vtable, &self.options, expression) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index f4862f6876a..9e18f030a6d 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -209,18 +209,12 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns an expression that evaluates to the validity of the result of this expression. /// - /// If a validity expression cannot be constructed, returns `None` and the expression will - /// be evaluated as normal before extracting the validity mask from the result. - /// - /// This is essentially a specialized form of a `reduce_parent` + /// This is essentially a specialized form of a `reduce_parent`. fn validity( &self, options: &Self::Options, expression: &Expression, - ) -> VortexResult> { - _ = (options, expression); - Ok(None) - } + ) -> VortexResult; /// Returns whether this expression itself is null-sensitive. Conservatively default to *true*. /// From 5b26b85320c4ced17ad511e4dda0099f945c3808 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 19 Jun 2026 15:12:27 -0400 Subject: [PATCH 17/19] Update tensor scalar validity signatures Signed-off-by: Nicholas Gates --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 4 ++-- vortex-tensor/src/scalar_fns/inner_product.rs | 4 ++-- vortex-tensor/src/scalar_fns/l2_denorm.rs | 4 ++-- vortex-tensor/src/scalar_fns/l2_norm.rs | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 82328a9b2d9..972f2e13166 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -178,12 +178,12 @@ impl ScalarFnVTable for CosineSimilarity { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { // The result is null if either input tensor is null. let lhs_validity = expression.child(0).validity()?; let rhs_validity = expression.child(1).validity()?; - Ok(Some(and(lhs_validity, rhs_validity))) + Ok(and(lhs_validity, rhs_validity)) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index cfea956f314..e11ac730c35 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -162,12 +162,12 @@ impl ScalarFnVTable for InnerProduct { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { // The result is null if either input tensor is null. let lhs_validity = expression.child(0).validity()?; let rhs_validity = expression.child(1).validity()?; - Ok(Some(and(lhs_validity, rhs_validity))) + Ok(and(lhs_validity, rhs_validity)) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 4f4159b4a21..c4ba7a3ba36 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -255,11 +255,11 @@ impl ScalarFnVTable for L2Denorm { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { let normalized_validity = expression.child(0).validity()?; let norms_validity = expression.child(1).validity()?; - Ok(Some(and(normalized_validity, norms_validity))) + Ok(and(normalized_validity, norms_validity)) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index aa85ab4c78e..ab7c387f595 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -182,9 +182,9 @@ impl ScalarFnVTable for L2Norm { &self, _options: &Self::Options, expression: &Expression, - ) -> VortexResult> { + ) -> VortexResult { // The result is null if the input tensor is null. - Ok(Some(expression.child(0).validity()?)) + expression.child(0).validity() } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { From aaf35d422eb33b2a133d0455f8d01e5ae572e9d8 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 19 Jun 2026 15:22:10 -0400 Subject: [PATCH 18/19] Fix scalar validity fallout in layout Signed-off-by: Nicholas Gates --- vortex-array/src/array/erased.rs | 3 ++- .../src/arrays/scalar_fn/vtable/operations.rs | 24 +++++++++++++++++++ vortex-layout/src/layouts/row_idx/expr.rs | 9 +++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 9b6d92014c9..9dad46fe0ae 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -41,6 +41,7 @@ use crate::arrays::DictArray; use crate::arrays::FilterArray; use crate::arrays::Null; use crate::arrays::Primitive; +use crate::arrays::ScalarFn; use crate::arrays::SliceArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; @@ -276,7 +277,7 @@ impl ArrayRef { /// Execute the array to extract a scalar at the given index. pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len()); - if self.dtype().is_nullable() && self.is_invalid(index, ctx)? { + if self.dtype().is_nullable() && !self.is::() && self.is_invalid(index, ctx)? { return Ok(Scalar::null(self.dtype().clone())); } let scalar = self.0.data.execute_scalar(self, index, ctx)?; diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 67de19d936f..782fdfda033 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -65,6 +65,7 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; + use crate::arrays::StructArray; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::assert_arrays_eq; use crate::scalar::Scalar; @@ -172,6 +173,29 @@ mod tests { Ok(()) } + #[test] + fn scalar_fn_scalar_at_handles_value_derived_validity() -> VortexResult<()> { + let child = StructArray::from_fields(&[( + "a", + PrimitiveArray::from_option_iter([Some(1i32), None]).into_array(), + )])? + .into_array(); + let expr = crate::expr::get_item("a", crate::expr::root()); + let array = ScalarFnArray::try_new(expr.scalar_fn().clone(), vec![child])?.into_array(); + + assert_eq!( + array.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?, + Scalar::primitive(1i32, true.into()) + ); + assert!( + array + .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())? + .is_null() + ); + + Ok(()) + } + #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { let lhs = buffer![1i32, 5, 3].into_array(); diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 95ddce3d762..122ee3ca9e8 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -8,6 +8,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::Expression; +use vortex_array::expr::lit; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -61,6 +62,14 @@ impl ScalarFnVTable for RowIdx { "RowIdxExpr should not be executed directly, use it in the context of a Vortex scan and it will be substituted for a row index array" ); } + + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } } pub fn row_idx() -> Expression { From ca773ed56db65e6e8b366016cd8c99b1cc1ed4f1 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 19 Jun 2026 15:34:10 -0400 Subject: [PATCH 19/19] Add remaining scalar function validity hooks Signed-off-by: Nicholas Gates --- vortex-geo/src/scalar_fn/distance.rs | 10 ++++++++++ vortex-row/src/encode.rs | 10 ++++++++++ vortex-row/src/size.rs | 10 ++++++++++ 3 files changed, 30 insertions(+) diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index 7f222cb763a..cf730a4cfb1 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -13,6 +13,8 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::lit; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -125,6 +127,14 @@ impl ScalarFnVTable for GeoDistance { } } } + + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } } /// Distance from each row of `points` to a constant `query` point, decoded once and broadcast. diff --git a/vortex-row/src/encode.rs b/vortex-row/src/encode.rs index 04feec89415..906d283abcf 100644 --- a/vortex-row/src/encode.rs +++ b/vortex-row/src/encode.rs @@ -18,6 +18,8 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::lit; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::ExecutionArgs; @@ -90,6 +92,14 @@ impl ScalarFnVTable for RowEncode { execute_row_encode(options, args, ctx) } + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { true } diff --git a/vortex-row/src/size.rs b/vortex-row/src/size.rs index 6636c4e9f34..b9f932dbf3d 100644 --- a/vortex-row/src/size.rs +++ b/vortex-row/src/size.rs @@ -18,6 +18,8 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::lit; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -242,6 +244,14 @@ impl ScalarFnVTable for RowSize { .into_array()) } + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult { + Ok(lit(true)) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { true }