vortex_array/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::type_name;
5use std::fmt::Debug;
6use std::marker::PhantomData;
7
8use vortex_error::VortexResult;
9use vortex_vector::Vector;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::matchers::MatchKey;
14use crate::matchers::Matcher;
15use crate::vtable::VTable;
16
17pub struct ParentKernelSet<V: VTable> {
18    kernels: &'static [&'static dyn DynParentKernel<V>],
19}
20
21impl<V: VTable> ParentKernelSet<V> {
22    /// Create a new parent kernel set with the given kernels.
23    ///
24    /// Use [`ParentKernelSet::lift`] to lift static rules into dynamic trait objects.
25    pub const fn new(kernels: &'static [&'static dyn DynParentKernel<V>]) -> Self {
26        Self { kernels }
27    }
28
29    /// Lift the given rule into a dynamic trait object.
30    pub const fn lift<K: ExecuteParentKernel<V>>(
31        kernel: &'static K,
32    ) -> &'static dyn DynParentKernel<V> {
33        // Assert that self is zero-sized
34        const {
35            assert!(
36                !(size_of::<K>() != 0),
37                "Rule must be zero-sized to be lifted"
38            );
39        }
40        unsafe { &*(kernel as *const K as *const ParentKernelAdapter<V, K>) }
41    }
42
43    /// Evaluate the parent kernels on the given child and parent arrays.
44    pub fn execute(
45        &self,
46        child: &V::Array,
47        parent: &ArrayRef,
48        child_idx: usize,
49        ctx: &mut ExecutionCtx,
50    ) -> VortexResult<Option<Vector>> {
51        for kernel in self.kernels.iter() {
52            if let MatchKey::Array(id) = kernel.parent_key()
53                && parent.encoding_id() != id
54            {
55                continue;
56            }
57            if let Some(reduced) = kernel.execute_parent(child, parent, child_idx, ctx)? {
58                return Ok(Some(reduced));
59            }
60        }
61        Ok(None)
62    }
63}
64
65pub trait ExecuteParentKernel<V: VTable>: Debug {
66    type Parent: Matcher;
67
68    /// Returns the matcher for the parent array
69    fn parent(&self) -> Self::Parent;
70
71    /// Attempt to execute the parent array fused with the child array.
72    fn execute_parent(
73        &self,
74        array: &V::Array,
75        parent: <Self::Parent as Matcher>::View<'_>,
76        child_idx: usize,
77        ctx: &mut ExecutionCtx,
78    ) -> VortexResult<Option<Vector>>;
79}
80
81pub trait DynParentKernel<V: VTable> {
82    fn parent_key(&self) -> MatchKey;
83
84    fn execute_parent(
85        &self,
86        child: &V::Array,
87        parent: &ArrayRef,
88        child_idx: usize,
89        ctx: &mut ExecutionCtx,
90    ) -> VortexResult<Option<Vector>>;
91}
92
93pub struct ParentKernelAdapter<V, K> {
94    kernel: K,
95    _phantom: PhantomData<V>,
96}
97
98impl<V: VTable, K: ExecuteParentKernel<V>> Debug for ParentKernelAdapter<V, K> {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        f.debug_struct("ParentKernelAdapter")
101            .field("parent", &type_name::<K::Parent>())
102            .field("kernel", &self.kernel)
103            .finish()
104    }
105}
106
107impl<V: VTable, R: ExecuteParentKernel<V>> DynParentKernel<V> for ParentKernelAdapter<V, R> {
108    fn parent_key(&self) -> MatchKey {
109        self.kernel.parent().key()
110    }
111
112    fn execute_parent(
113        &self,
114        child: &V::Array,
115        parent: &ArrayRef,
116        child_idx: usize,
117        ctx: &mut ExecutionCtx,
118    ) -> VortexResult<Option<Vector>> {
119        let Some(parent_view) = self.kernel.parent().try_match(parent) else {
120            return Ok(None);
121        };
122        self.kernel
123            .execute_parent(child, parent_view, child_idx, ctx)
124    }
125}
vortex_array/kernel.rs

vortex_array/
kernel.rs
