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

vortex_array/
kernel.rs
