vortex_array/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Parent kernels: child-driven fused execution of parent arrays.
5//!
6//! A parent kernel allows a child encoding to provide a specialized execution path for its
7//! parent array. This is Layer 3 of the [execution model](https://docs.vortex.dev/developer-guide/internals/execution).
8//!
9//! For example, a `RunEndArray` child of a `SliceArray` can perform a binary search on its
10//! run ends rather than decoding the entire array and slicing the result.
11//!
12//! Encodings declare their parent kernels by implementing [`ExecuteParentKernel`] and
13//! registering them in a [`ParentKernelSet`]. Each kernel specifies which parent types it
14//! handles via a [`Matcher`].
15
16use std::any::type_name;
17use std::fmt::Debug;
18use std::marker::PhantomData;
19
20use vortex_error::VortexResult;
21
22use crate::ArrayRef;
23use crate::ExecutionCtx;
24use crate::matcher::Matcher;
25use crate::vtable::VTable;
26
27/// A collection of [`ExecuteParentKernel`]s registered for a specific child encoding.
28///
29/// During execution, the scheduler iterates over each child's `ParentKernelSet` looking for
30/// a kernel whose [`Matcher`] matches the parent array type. The first matching kernel that
31/// returns `Some` wins.
32pub struct ParentKernelSet<V: VTable> {
33    kernels: &'static [&'static dyn DynParentKernel<V>],
34}
35
36impl<V: VTable> ParentKernelSet<V> {
37    /// Create a new parent kernel set with the given kernels.
38    ///
39    /// Use [`ParentKernelSet::lift`] to lift static rules into dynamic trait objects.
40    pub const fn new(kernels: &'static [&'static dyn DynParentKernel<V>]) -> Self {
41        Self { kernels }
42    }
43
44    /// Lift the given rule into a dynamic trait object.
45    pub const fn lift<K: ExecuteParentKernel<V>>(
46        kernel: &'static K,
47    ) -> &'static dyn DynParentKernel<V> {
48        // Assert that self is zero-sized
49        const {
50            assert!(
51                !(size_of::<K>() != 0),
52                "Rule must be zero-sized to be lifted"
53            );
54        }
55        unsafe { &*(kernel as *const K as *const ParentKernelAdapter<V, K>) }
56    }
57
58    /// Evaluate the parent kernels on the given child and parent arrays.
59    pub fn execute(
60        &self,
61        child: &V::Array,
62        parent: &ArrayRef,
63        child_idx: usize,
64        ctx: &mut ExecutionCtx,
65    ) -> VortexResult<Option<ArrayRef>> {
66        for kernel in self.kernels.iter() {
67            if !kernel.matches(parent) {
68                continue;
69            }
70            if let Some(reduced) = kernel.execute_parent(child, parent, child_idx, ctx)? {
71                return Ok(Some(reduced));
72            }
73        }
74        Ok(None)
75    }
76}
77
78/// A kernel that allows a child encoding `V` to execute its parent array in a fused manner.
79///
80/// This is the typed trait that encoding authors implement. The associated `Parent` type
81/// specifies which parent array types this kernel can handle. When the parent matches,
82/// [`execute_parent`](Self::execute_parent) is called with the strongly-typed child and parent views.
83///
84/// Unlike reduce rules, parent kernels may read buffers and perform real computation.
85///
86/// Return `Ok(None)` to decline handling (the scheduler will try the next kernel or fall
87/// through to the encoding's own `execute`).
88pub trait ExecuteParentKernel<V: VTable>: Debug + Send + Sync + 'static {
89    /// The parent array type this kernel handles.
90    type Parent: Matcher;
91
92    /// Attempt to execute the parent array fused with the child array.
93    fn execute_parent(
94        &self,
95        array: &V::Array,
96        parent: <Self::Parent as Matcher>::Match<'_>,
97        child_idx: usize,
98        ctx: &mut ExecutionCtx,
99    ) -> VortexResult<Option<ArrayRef>>;
100}
101
102/// Type-erased version of [`ExecuteParentKernel`] used for dynamic dispatch within
103/// [`ParentKernelSet`].
104pub trait DynParentKernel<V: VTable>: Send + Sync {
105    /// Returns `true` if this kernel's parent [`Matcher`] matches the given parent array.
106    fn matches(&self, parent: &ArrayRef) -> bool;
107
108    /// Attempt to execute the parent array fused with the child array.
109    fn execute_parent(
110        &self,
111        child: &V::Array,
112        parent: &ArrayRef,
113        child_idx: usize,
114        ctx: &mut ExecutionCtx,
115    ) -> VortexResult<Option<ArrayRef>>;
116}
117
118/// Bridges a concrete [`ExecuteParentKernel<V, K>`] to the type-erased [`DynParentKernel<V>`]
119/// trait. Created by [`ParentKernelSet::lift`].
120pub struct ParentKernelAdapter<V, K> {
121    kernel: K,
122    _phantom: PhantomData<V>,
123}
124
125impl<V: VTable, K: ExecuteParentKernel<V>> Debug for ParentKernelAdapter<V, K> {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        f.debug_struct("ParentKernelAdapter")
128            .field("parent", &type_name::<K::Parent>())
129            .field("kernel", &self.kernel)
130            .finish()
131    }
132}
133
134impl<V: VTable, K: ExecuteParentKernel<V>> DynParentKernel<V> for ParentKernelAdapter<V, K> {
135    fn matches(&self, parent: &ArrayRef) -> bool {
136        K::Parent::matches(parent)
137    }
138
139    fn execute_parent(
140        &self,
141        child: &V::Array,
142        parent: &ArrayRef,
143        child_idx: usize,
144        ctx: &mut ExecutionCtx,
145    ) -> VortexResult<Option<ArrayRef>> {
146        let Some(parent_view) = K::Parent::try_match(parent) else {
147            return Ok(None);
148        };
149        self.kernel
150            .execute_parent(child, parent_view, child_idx, ctx)
151    }
152}
vortex_array/kernel.rs

vortex_array/
kernel.rs
