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

vortex_array/
kernel.rs
