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