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/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 7801316445e..678da8642a7 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/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::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::BooleanExecuteAdaptor; +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(&BooleanExecuteAdaptor(ByteBool)), - ParentKernelSet::lift(&CastExecuteAdaptor(ByteBool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), ByteBool, BooleanExecuteAdaptor(ByteBool)); + 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/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/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/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/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..92ff63f4408 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -56,6 +56,7 @@ mod test { use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; + 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() @@ -154,12 +155,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..458b5d7563f 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( @@ -137,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; @@ -145,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/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/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/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..5624437d461 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; @@ -232,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)); @@ -356,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 { @@ -807,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/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/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/runend/benches/run_end_take.rs b/encodings/runend/benches/run_end_take.rs index 1c9d5e7ba7d..f0dccccfc48 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/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/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 60644e2bced..83eb811e682 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -116,7 +116,6 @@ 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; @@ -125,11 +124,12 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; + use crate::tests::SESSION; fn ree_array() -> RunEndArray { 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 +139,7 @@ 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 mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( filtered, RunEnd::new( @@ -157,7 +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 mut ctx = LEGACY_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/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..8169f5d8dfc 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( @@ -59,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() { 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/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/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index dd84123851a..17b4241f293 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(&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(&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(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(fill_elements, ctx) .vortex_expect("Failed to append sparse fixed-size-list fill value"); } } else { 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/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/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 cbcf1727ebc..c19d00d5118 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/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::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::binary::Binary; use crate::scalar_fn::fns::binary::BooleanExecuteAdaptor; +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(&BooleanExecuteAdaptor(Bool)), - 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(Binary.id(), Bool, BooleanExecuteAdaptor(Bool)); + 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..eaa39462894 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; @@ -24,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; @@ -43,6 +44,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 @@ -170,17 +175,21 @@ impl VTable for Bool { SLOT_NAMES[idx].to_string() } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn execute_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, + builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + ) -> 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)) } fn reduce_parent( 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..cd2ef72a7ee 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; @@ -26,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; @@ -38,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; @@ -54,6 +55,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; @@ -196,20 +201,38 @@ impl VTable for Dict { )?)) } - fn reduce_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + 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 execute_parent( + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, child_idx: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + PARENT_RULES.evaluate(array, parent, child_idx) } } 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..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,12 +19,28 @@ 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::kernel::ExecuteParentKernel; use crate::validity::Validity; +fn execute_parent( + child: &crate::ArrayRef, + parent: &crate::ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + TakeExecuteAdaptor(Filter).execute_parent( + child.as_::(), + parent.as_::(), + child_idx, + ctx, + ) +} + #[test] fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { let filter = FilterArray::new( @@ -42,8 +60,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 +88,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 +109,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 +129,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 +207,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 +231,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 +248,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 +270,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 +337,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 +359,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 +385,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 +478,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..d9d778c34bb 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; @@ -15,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; @@ -39,6 +41,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); @@ -184,21 +190,27 @@ impl VTable for Primitive { Ok(ExecutionResult::done(array)) } - fn reduce_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + 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 execute_parent( + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, child_idx: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + RULES.evaluate(array, parent, child_idx) } } 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..a9ffe5a5859 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; @@ -29,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; @@ -40,6 +41,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,13 +222,17 @@ impl VTable for VarBinView { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, + builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + ) -> 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 { 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/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/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7a3c3bb068a..2e7aced758b 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; @@ -83,7 +84,11 @@ 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<()> { + 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 {:?}", @@ -97,7 +102,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(()) @@ -997,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 @@ -1018,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 @@ -1048,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 4c350f1888d..fef03662181 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; @@ -117,7 +118,11 @@ 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<()> { + 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 {:?}", @@ -135,7 +140,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( @@ -773,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 @@ -794,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 @@ -829,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() + ); } } 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) { diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 2bd5dfd5b37..f90400c3412 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -15,6 +15,7 @@ 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; @@ -41,6 +42,8 @@ use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ParentExecutionKernels; +use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; @@ -138,12 +141,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` @@ -160,6 +167,8 @@ impl ArrayRef { let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); + let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); + let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); for _ in 0..max_iterations { @@ -196,8 +205,15 @@ 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) = { + execute_parent_for_child( + &frame.parent_array, + ¤t_array, + frame.slot_idx, + kernels, + ctx, + )? + } { ctx.log(format_args!( "execute_parent (stack) rewrote {} -> {}", @@ -211,7 +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, ctx)? + && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { ctx.log(format_args!( "execute_parent rewrote {} -> {}", @@ -303,6 +319,7 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, + execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, #[cfg(debug_assertions)] @@ -312,8 +329,13 @@ pub struct ExecutionCtx { 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) + .unwrap_or_default(); Self { session, + execute_parent_kernels, #[cfg(debug_assertions)] id: { static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0); @@ -422,8 +444,8 @@ impl Executable for ArrayRef { } } - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); + let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); + let kernels = execute_parent_kernels.as_ref(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; @@ -540,28 +562,27 @@ fn execute_parent_for_child( parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, - kernels: Option<&ArrayKernels>, + kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(kernels) = kernels - && let Some(plugins) = - kernels.find_execute_parent(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(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. -fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); - +fn try_execute_parent( + array: &ArrayRef, + kernels: &ParentExecutionKernels, + ctx: &mut ExecutionCtx, +) -> VortexResult> { for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = 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..f2f89498b97 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; @@ -34,14 +34,16 @@ 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; 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,9 +92,69 @@ 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>; +} + +pub(crate) type ExecuteParentKernelRef = Arc; + +pub(crate) type ParentExecutionKernels = HashMap>; + +#[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); +pub(crate) struct ExecuteParentFnId(u64); impl From for ExecuteParentFnId { fn from(id: u64) -> Self { @@ -113,14 +175,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 +203,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,28 +230,66 @@ 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)) } + + /// Return the currently published execute-parent kernel snapshot. + pub(crate) fn execute_parent_snapshot(&self) -> Arc { + self.execute_parent.snapshot() + } } 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 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. 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/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(); 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 = {