Optimize arrays only with a session; keep constructors lazy

encodings/alp/src/alp/array.rs

@@ -51,7 +51,6 @@ use vortex_session::registry::CachedId;
 use crate::ALPFloat;
 use crate::alp::Exponents;
 use crate::alp::decompress::execute_decompress;
-use crate::alp::rules::RULES;
 
 /// A [`ALP`]-encoded Vortex array.
 pub type ALPArray = Array<ALP>;
@@ -203,14 +202,6 @@ impl VTable for ALP {
             execute_decompress(array, ctx)?.into_array(),
         ))
     }
-
-    fn reduce_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-    ) -> VortexResult<Option<ArrayRef>> {
-        RULES.evaluate(array, parent, child_idx)
-    }
 }
 
 #[array_slots(ALP)]

encodings/alp/src/alp/compute/filter.rs

@@ -54,7 +54,11 @@ mod test {
 
     use crate::alp_encode;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[rstest]
     #[case(buffer![1.23f32, 4.56, 7.89, 10.11, 12.13].into_array())]

encodings/alp/src/alp/compute/mask.rs

@@ -65,7 +65,11 @@ mod test {
     use crate::alp::array::ALPArrayExt;
     use crate::alp_encode;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[rstest]
     #[case(buffer![10.5f32, 20.5, 30.5, 40.5, 50.5].into_array())]

encodings/alp/src/alp/compute/take.rs

@@ -44,7 +44,11 @@ mod test {
 
     use crate::alp_encode;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[rstest]
     #[case(buffer![1.23f32, 4.56, 7.89, 10.11, 12.13].into_array())]

encodings/alp/src/alp/mod.rs

@@ -22,6 +22,7 @@ pub(crate) use plugin::ALPPatchedPlugin;
 
 pub(crate) fn initialize(session: &VortexSession) {
     rules::initialize(session);
+    rules::register_parent_rules(session);
 }
 
 #[cfg(test)]

encodings/alp/src/alp/rules.rs

@@ -9,12 +9,14 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor;
 use vortex_array::arrays::filter::FilterExecuteAdaptor;
 use vortex_array::arrays::slice::SliceExecuteAdaptor;
 use vortex_array::optimizer::kernels::ArrayKernelsExt;
-use vortex_array::optimizer::rules::ParentRuleSet;
 use vortex_array::scalar_fn::ScalarFnVTable;
+use vortex_array::scalar_fn::fns::between::Between;
 use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor;
 use vortex_array::scalar_fn::fns::binary::Binary;
 use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor;
+use vortex_array::scalar_fn::fns::cast::Cast;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
+use vortex_array::scalar_fn::fns::mask::Mask as MaskFn;
 use vortex_array::scalar_fn::fns::mask::Mask;
 use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor;
 use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor;
@@ -31,8 +33,10 @@ pub(super) fn initialize(session: &VortexSession) {
     kernels.register_execute_parent_kernel(Dict.id(), ALP, TakeExecuteAdaptor(ALP));
 }
 
-pub(super) const RULES: ParentRuleSet<ALP> = ParentRuleSet::new(&[
-    ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)),
-    ParentRuleSet::lift(&CastReduceAdaptor(ALP)),
-    ParentRuleSet::lift(&MaskReduceAdaptor(ALP)),
-]);
+/// Registers the parent-reduce rules for [`ALP`] children in the session kernels.
+pub(crate) fn register_parent_rules(session: &VortexSession) {
+    let kernels = session.kernels();
+    kernels.register_reduce_parent_rule(Between.id(), ALP, BetweenReduceAdaptor(ALP));
+    kernels.register_reduce_parent_rule(Cast.id(), ALP, CastReduceAdaptor(ALP));
+    kernels.register_reduce_parent_rule(MaskFn.id(), ALP, MaskReduceAdaptor(ALP));
+}

encodings/alp/src/alp_rd/array.rs

@@ -55,7 +55,6 @@ use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
 use crate::alp_rd::check_left_parts_codes;
-use crate::alp_rd::rules::RULES;
 use crate::alp_rd_decode;
 
 /// A [`ALPRD`]-encoded Vortex array.
@@ -335,14 +334,6 @@ impl VTable for ALPRD {
 
         Ok(ExecutionResult::done(decoded_array.into_array()))
     }
-
-    fn reduce_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-    ) -> VortexResult<Option<ArrayRef>> {
-        RULES.evaluate(array, parent, child_idx)
-    }
 }
 
 #[array_slots(ALPRD)]

encodings/alp/src/alp_rd/compute/cast.rs

@@ -64,7 +64,11 @@ mod tests {
     use crate::RDEncoder;
     use crate::RDEncoderExt;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[test]
     fn test_cast_alprd_f32_to_f64() {

encodings/alp/src/alp_rd/mod.rs

@@ -46,6 +46,7 @@ mod slice;
 
 pub(crate) fn initialize(session: &VortexSession) {
     kernel::initialize(session);
+    rules::register_parent_rules(session);
 }
 
 /// Encoding of Vortex arrays with an ALP-RD [`RDEncoder`].

encodings/alp/src/alp_rd/rules.rs

@@ -1,13 +1,19 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_array::optimizer::kernels::ArrayKernelsExt;
+use vortex_array::scalar_fn::ScalarFnVTable;
+use vortex_array::scalar_fn::fns::cast::Cast;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
+use vortex_array::scalar_fn::fns::mask::Mask as MaskFn;
 use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor;
+use vortex_session::VortexSession;
 
 use crate::alp_rd::ALPRD;
 
-pub(crate) static RULES: ParentRuleSet<ALPRD> = ParentRuleSet::new(&[
-    ParentRuleSet::lift(&CastReduceAdaptor(ALPRD)),
-    ParentRuleSet::lift(&MaskReduceAdaptor(ALPRD)),
-]);
+/// Registers the parent-reduce rules for [`ALPRD`] children in the session kernels.
+pub(crate) fn register_parent_rules(session: &VortexSession) {
+    let kernels = session.kernels();
+    kernels.register_reduce_parent_rule(Cast.id(), ALPRD, CastReduceAdaptor(ALPRD));
+    kernels.register_reduce_parent_rule(MaskFn.id(), ALPRD, MaskReduceAdaptor(ALPRD));
+}

encodings/bytebool/src/array.rs

@@ -163,14 +163,6 @@ impl VTable for ByteBool {
         ByteBoolSlots::NAMES[idx].to_string()
     }
 
-    fn reduce_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-    ) -> VortexResult<Option<ArrayRef>> {
-        crate::rules::RULES.evaluate(array, parent, child_idx)
-    }
-
     fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         // convert truthy values to set/unset bits
         let boolean_buffer = BitBufferMut::from(array.truthy_bytes()).freeze();

encodings/bytebool/src/lib.rs

@@ -55,4 +55,5 @@ mod slice;
 pub fn initialize(session: &VortexSession) {
     session.arrays().register(ByteBool);
     kernel::initialize(session);
+    rules::register_parent_rules(session);
 }

encodings/bytebool/src/rules.rs

@@ -1,15 +1,23 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_array::ArrayVTable;
+use vortex_array::arrays::Slice;
 use vortex_array::arrays::slice::SliceReduceAdaptor;
-use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_array::optimizer::kernels::ArrayKernelsExt;
+use vortex_array::scalar_fn::ScalarFnVTable;
+use vortex_array::scalar_fn::fns::cast::Cast;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
+use vortex_array::scalar_fn::fns::mask::Mask as MaskFn;
 use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor;
+use vortex_session::VortexSession;
 
 use crate::ByteBool;
 
-pub(crate) static RULES: ParentRuleSet<ByteBool> = ParentRuleSet::new(&[
-    ParentRuleSet::lift(&CastReduceAdaptor(ByteBool)),
-    ParentRuleSet::lift(&MaskReduceAdaptor(ByteBool)),
-    ParentRuleSet::lift(&SliceReduceAdaptor(ByteBool)),
-]);
+/// Registers the parent-reduce rules for [`ByteBool`] children in the session kernels.
+pub(crate) fn register_parent_rules(session: &VortexSession) {
+    let kernels = session.kernels();
+    kernels.register_reduce_parent_rule(Cast.id(), ByteBool, CastReduceAdaptor(ByteBool));
+    kernels.register_reduce_parent_rule(MaskFn.id(), ByteBool, MaskReduceAdaptor(ByteBool));
+    kernels.register_reduce_parent_rule(Slice.id(), ByteBool, SliceReduceAdaptor(ByteBool));
+}

encodings/datetime-parts/src/array.rs

@@ -42,7 +42,6 @@ use vortex_session::registry::CachedId;
 
 use crate::TemporalParts;
 use crate::canonical::decode_to_temporal;
-use crate::compute::rules::PARENT_RULES;
 use crate::split_temporal;
 
 /// A [`DateTimeParts`]-encoded Vortex array.
@@ -205,14 +204,6 @@ impl VTable for DateTimeParts {
             decode_to_temporal(parts, &dtype, ctx)?.into_array(),
         ))
     }
-
-    fn reduce_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-    ) -> VortexResult<Option<ArrayRef>> {
-        PARENT_RULES.evaluate(array, parent, child_idx)
-    }
 }
 
 #[array_slots(DateTimeParts)]

encodings/datetime-parts/src/compress.rs

@@ -186,7 +186,11 @@ mod tests {
     use crate::TemporalParts;
     use crate::split_temporal;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[rstest]
     #[case(Validity::NonNullable)]

encodings/datetime-parts/src/compute/cast.rs

@@ -53,7 +53,11 @@ mod tests {
     use crate::DateTimeParts;
     use crate::DateTimePartsArray;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     fn date_time_array(validity: Validity) -> ArrayRef {
         DateTimeParts::try_from_temporal(

encodings/datetime-parts/src/compute/compare.rs

@@ -311,7 +311,11 @@ mod test {
     use crate::DateTimeParts;
     use crate::DateTimePartsArray;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     fn dtp_array_from_timestamp<T: IntegerPType>(
         value: T,

encodings/datetime-parts/src/compute/filter.rs

@@ -39,7 +39,11 @@ mod test {
 
     use crate::DateTimeParts;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     #[test]
     fn test_filter_datetime_parts() {

encodings/datetime-parts/src/compute/mod.rs

@@ -28,7 +28,11 @@ mod tests {
     use crate::DateTimeParts;
     use crate::DateTimePartsArray;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty();
+        crate::initialize(&session);
+        session
+    });
 
     fn dtp_from_temporal(temporal: TemporalArray) -> DateTimePartsArray {
         DateTimeParts::try_from_temporal(temporal, &mut SESSION.create_execution_ctx()).unwrap()