Merge origin/develop into aggregate-zoned-stats

Signed-off-by: "Nicholas Gates" <nick@nickgates.com>

Author: gatesn

docs/developer-guide/internals/execution.md

@@ -115,21 +115,28 @@ Each child is given the opportunity to execute its parent in a fused manner via
 `ExecuteParentKernel`. Unlike reduce rules, parent kernels may read buffers and perform real
 computation.
 
-An encoding declares its parent kernels in a `ParentKernelSet`, specifying which parent types
-each kernel handles via a `Matcher`:
+An encoding declares parent kernels by implementing `ExecuteParentKernel` and registering them
+with the session-scoped kernel registry, specifying which parent types each kernel handles via a
+`Matcher`:
 
 ```rust
 pub trait ExecuteParentKernel<V: VTable> {
     type Parent: Matcher;  // which parent types this kernel handles
 
     fn execute_parent(
         &self,
-        array: &V::Array,                          // the child
+        array: ArrayView<'_, V>,                   // the child
         parent: <Self::Parent as Matcher>::Match<'_>, // the matched parent
         child_idx: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>>;
 }
+
+pub fn initialize(session: &VortexSession) {
+    session
+        .kernels()
+        .register_execute_parent_kernel(parent_id, Child, Kernel);
+}
 ```
 
 Examples:
@@ -195,17 +202,23 @@ execute_until<M>(root):
     │     restore builder, │
     │     loop             ▼
     │         ┌────────────────────────────────────────────┐
-    │         │ Step 2a: current_array.execute_parent(     │
-    │         │            stack.top.parent_array )        │
-    │         │ child looks UP at the suspended parent     │
+    │         │ Step 2a: execute_parent_for_child(         │
+    │         │            stack.top.parent_array,         │
+    │         │            current_array )                 │
+    │         │ lookup session kernels by key:             │
+    │         │   (parent.encoding_id(),                   │
+    │         │    child.encoding_id())                    │
     │         ├────────────┬───────────────────────────────┘
     │         │ Some       │ None
     │         │            │
     │         │            ▼
     │         │  ┌─────────────────────────────────────────┐
-    │         │  │ Step 2b: each child.execute_parent(     │
-    │         │  │            current_array )              │
-    │         │  │ children look UP at current_array       │
+    │         │  │ Step 2b: for each child:                │
+    │         │  │   execute_parent_for_child(             │
+    │         │  │     current_array, child )              │
+    │         │  │ lookup session kernels by key:          │
+    │         │  │   (parent.encoding_id(),                │
+    │         │  │    child.encoding_id())                 │
     │         │  ├──────────┬──────────────────────────────┘
     │         │  │ Some     │ None
     │         │  │          │
@@ -236,6 +249,9 @@ Step 2a and Step 2b are skipped while `current_builder` is active. `AppendChild`
 consumes `current_array`: some slots already live in the builder, so a parent rewrite would
 observe inconsistent state and could discard accumulated builder data.
 
+Both Step 2a and Step 2b route through `execute_parent_for_child`, which looks up
+`(parent.encoding_id(), child.encoding_id())` in the session kernel snapshot.
+
 ## Incremental Execution
 
 Execution is incremental: each call to `execute` moves the array one step closer to canonical
@@ -256,12 +272,12 @@ codes through the slice, missing the Dict-RLE optimization entirely. Incremental
 avoids this:
 
 1. First iteration: the slice `execute` returns `ExecuteSlot` for its `RunEndArray` child.
-   Once that child is in focus, Step 2a gives it a chance to rewrite the suspended slice
-   parent before the child is forced toward canonical form.
+   Once that child is in focus, Step 2a looks up a registered parent kernel for the
+   `(slice, runend)` pair before the child is forced toward canonical form.
 
 2. Second iteration: the `RunEndArray` codes child now matches the Dict-RLE pattern. Its
-   `execute_parent` provides a fused kernel that expands runs while performing dictionary
-   lookups in a single pass, returning the canonical array directly.
+   registered execute-parent kernel expands runs while performing dictionary lookups in a single
+   pass, returning the canonical array directly.
 
 ## Walkthrough: Executing a RunEnd-Encoded Array
 

encodings/alp/benches/alp_compress.rs

@@ -50,7 +50,11 @@ const BENCH_ARGS: &[(usize, f64, f64)] = &[
     (10_000, 0.1, 1.0),
 ];
 
-static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_alp::initialize(&session);
+    session
+});
 
 #[divan::bench(types = [f32, f64], args = BENCH_ARGS)]
 fn compress_alp<T: ALPFloat + NativePType>(bencher: Bencher, args: (usize, f64, f64)) {

encodings/alp/src/alp/array.rs

@@ -48,7 +48,6 @@ use vortex_session::registry::CachedId;
 use crate::ALPFloat;
 use crate::alp::Exponents;
 use crate::alp::decompress::execute_decompress;
-use crate::alp::rules::PARENT_KERNELS;
 use crate::alp::rules::RULES;
 
 /// A [`ALP`]-encoded Vortex array.
@@ -188,15 +187,6 @@ impl VTable for ALP {
     ) -> VortexResult<Option<ArrayRef>> {
         RULES.evaluate(array, parent, child_idx)
     }
-
-    fn execute_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<Option<ArrayRef>> {
-        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
-    }
 }
 
 #[array_slots(ALP)]

encodings/alp/src/alp/mod.rs

@@ -61,9 +61,14 @@ use vortex_array::dtype::NativePType;
 use vortex_array::scalar::PValue;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
+use vortex_session::VortexSession;
 
 const SAMPLE_SIZE: usize = 32;
 
+pub(crate) fn initialize(session: &VortexSession) {
+    rules::initialize(session);
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub struct Exponents {
     pub e: u8,

encodings/alp/src/alp/rules.rs

@@ -1,26 +1,35 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_array::ArrayVTable;
+use vortex_array::arrays::Dict;
+use vortex_array::arrays::Filter;
+use vortex_array::arrays::Slice;
 use vortex_array::arrays::dict::TakeExecuteAdaptor;
 use vortex_array::arrays::filter::FilterExecuteAdaptor;
 use vortex_array::arrays::slice::SliceExecuteAdaptor;
-use vortex_array::kernel::ParentKernelSet;
+use vortex_array::optimizer::kernels::ArrayKernelsExt;
 use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_array::scalar_fn::ScalarFnVTable;
 use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor;
+use vortex_array::scalar_fn::fns::binary::Binary;
 use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
+use vortex_array::scalar_fn::fns::mask::Mask;
 use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor;
 use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor;
+use vortex_session::VortexSession;
 
 use crate::ALP;
 
-pub(super) const PARENT_KERNELS: ParentKernelSet<ALP> = ParentKernelSet::new(&[
-    ParentKernelSet::lift(&CompareExecuteAdaptor(ALP)),
-    ParentKernelSet::lift(&FilterExecuteAdaptor(ALP)),
-    ParentKernelSet::lift(&MaskExecuteAdaptor(ALP)),
-    ParentKernelSet::lift(&SliceExecuteAdaptor(ALP)),
-    ParentKernelSet::lift(&TakeExecuteAdaptor(ALP)),
-]);
+pub(super) fn initialize(session: &VortexSession) {
+    let kernels = session.kernels();
+    kernels.register_execute_parent_kernel(Binary.id(), ALP, CompareExecuteAdaptor(ALP));
+    kernels.register_execute_parent_kernel(Filter.id(), ALP, FilterExecuteAdaptor(ALP));
+    kernels.register_execute_parent_kernel(Mask.id(), ALP, MaskExecuteAdaptor(ALP));
+    kernels.register_execute_parent_kernel(Slice.id(), ALP, SliceExecuteAdaptor(ALP));
+    kernels.register_execute_parent_kernel(Dict.id(), ALP, TakeExecuteAdaptor(ALP));
+}
 
 pub(super) const RULES: ParentRuleSet<ALP> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)),

encodings/alp/src/alp_rd/array.rs

@@ -53,7 +53,6 @@ use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::alp_rd::kernel::PARENT_KERNELS;
 use crate::alp_rd::rules::RULES;
 use crate::alp_rd_decode;
 
@@ -307,15 +306,6 @@ impl VTable for ALPRD {
     ) -> VortexResult<Option<ArrayRef>> {
         RULES.evaluate(array, parent, child_idx)
     }
-
-    fn execute_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<Option<ArrayRef>> {
-        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
-    }
 }
 
 /// The left (most significant) parts of the real-double encoded values.

encodings/alp/src/alp_rd/kernel.rs

@@ -1,15 +1,21 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_array::ArrayVTable;
+use vortex_array::arrays::Dict;
+use vortex_array::arrays::Filter;
+use vortex_array::arrays::Slice;
 use vortex_array::arrays::dict::TakeExecuteAdaptor;
 use vortex_array::arrays::filter::FilterExecuteAdaptor;
 use vortex_array::arrays::slice::SliceExecuteAdaptor;
-use vortex_array::kernel::ParentKernelSet;
+use vortex_array::optimizer::kernels::ArrayKernelsExt;
+use vortex_session::VortexSession;
 
 use crate::alp_rd::ALPRD;
 
-pub(crate) static PARENT_KERNELS: ParentKernelSet<ALPRD> = ParentKernelSet::new(&[
-    ParentKernelSet::lift(&SliceExecuteAdaptor(ALPRD)),
-    ParentKernelSet::lift(&FilterExecuteAdaptor(ALPRD)),
-    ParentKernelSet::lift(&TakeExecuteAdaptor(ALPRD)),
-]);
+pub(crate) fn initialize(session: &VortexSession) {
+    let kernels = session.kernels();
+    kernels.register_execute_parent_kernel(Slice.id(), ALPRD, SliceExecuteAdaptor(ALPRD));
+    kernels.register_execute_parent_kernel(Filter.id(), ALPRD, FilterExecuteAdaptor(ALPRD));
+    kernels.register_execute_parent_kernel(Dict.id(), ALPRD, TakeExecuteAdaptor(ALPRD));
+}

encodings/alp/src/alp_rd/mod.rs

@@ -36,6 +36,7 @@ use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_panic;
+use vortex_session::VortexSession;
 use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::match_each_alp_float_ptype;
@@ -55,6 +56,10 @@ const CUT_LIMIT: usize = 16;
 
 const MAX_DICT_SIZE: u8 = 8;
 
+pub(crate) fn initialize(session: &VortexSession) {
+    kernel::initialize(session);
+}
+
 mod private {
     pub trait Sealed {}
 

encodings/alp/src/lib.rs

@@ -39,6 +39,8 @@ pub fn initialize(session: &VortexSession) {
         session.arrays().register(ALP);
     }
     session.arrays().register(ALPRD);
+    alp::initialize(session);
+    alp_rd::initialize(session);
 
     // Register the ALP-specific NaN count aggregate kernel.
     session.aggregate_fns().register_aggregate_kernel(

encodings/bytebool/src/array.rs

@@ -39,8 +39,6 @@ use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::kernel::PARENT_KERNELS;
-
 /// A [`ByteBool`]-encoded Vortex array.
 pub type ByteBoolArray = Array<ByteBool>;
 
@@ -157,15 +155,6 @@ impl VTable for ByteBool {
             BoolArray::new(boolean_buffer, validity).into_array(),
         ))
     }
-
-    fn execute_parent(
-        array: ArrayView<'_, Self>,
-        parent: &ArrayRef,
-        child_idx: usize,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<Option<ArrayRef>> {
-        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
-    }
 }
 
 /// The validity bitmap indicating which elements are non-null.