vortex_array/arrays/slice/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod array;
5mod rules;
6mod slice_;
7mod vtable;
8
9use std::ops::Range;
10
11pub use array::*;
12use vortex_error::VortexResult;
13pub use vtable::*;
14
15use crate::ArrayRef;
16use crate::Canonical;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::kernel::ExecuteParentKernel;
20use crate::matcher::Matcher;
21use crate::optimizer::rules::ArrayParentReduceRule;
22use crate::vtable::VTable;
23
24pub trait SliceReduce: VTable {
25    /// Slice an array with the provided range without reading buffers.
26    ///
27    /// This trait is for slice implementations that can operate purely on array metadata and
28    /// structure without needing to read or execute on the underlying buffers. Implementations
29    /// should return `None` if slicing requires buffer access.
30    ///
31    /// # Preconditions
32    ///
33    /// The range is guaranteed to be within bounds of the array (i.e., `range.end <= array.len()`).
34    ///
35    /// Additionally, the range is guaranteed to be non-empty (i.e., `range.start < range.end`).
36    fn slice(array: &Self::Array, range: Range<usize>) -> VortexResult<Option<ArrayRef>>;
37}
38
39pub trait SliceKernel: VTable {
40    /// Slice an array with the provided range, potentially reading buffers.
41    ///
42    /// Unlike [`SliceReduce`], this trait is for slice implementations that may need to read
43    /// and execute on the underlying buffers to produce the sliced result.
44    ///
45    /// # Preconditions
46    ///
47    /// The range is guaranteed to be within bounds of the array (i.e., `range.end <= array.len()`).
48    ///
49    /// Additionally, the range is guaranteed to be non-empty (i.e., `range.start < range.end`).
50    fn slice(
51        array: &Self::Array,
52        range: Range<usize>,
53        ctx: &mut ExecutionCtx,
54    ) -> VortexResult<Option<ArrayRef>>;
55}
56
57fn precondition<V: VTable>(array: &V::Array, range: &Range<usize>) -> Option<ArrayRef> {
58    if range.start == 0 && range.end == array.len() {
59        return Some(array.to_array());
60    };
61    if range.start == range.end {
62        return Some(Canonical::empty(array.dtype()).into_array());
63    }
64    None
65}
66
67#[derive(Default, Debug)]
68pub struct SliceReduceAdaptor<V>(pub V);
69
70impl<V> ArrayParentReduceRule<V> for SliceReduceAdaptor<V>
71where
72    V: SliceReduce,
73{
74    type Parent = SliceVTable;
75
76    fn reduce_parent(
77        &self,
78        array: &V::Array,
79        parent: <Self::Parent as Matcher>::Match<'_>,
80        child_idx: usize,
81    ) -> VortexResult<Option<ArrayRef>> {
82        assert_eq!(child_idx, 0);
83        if let Some(result) = precondition::<V>(array, &parent.range) {
84            return Ok(Some(result));
85        }
86        <V as SliceReduce>::slice(array, parent.range.clone())
87    }
88}
89
90#[derive(Default, Debug)]
91pub struct SliceExecuteAdaptor<V>(pub V);
92
93impl<V> ExecuteParentKernel<V> for SliceExecuteAdaptor<V>
94where
95    V: SliceKernel,
96{
97    type Parent = SliceVTable;
98
99    fn execute_parent(
100        &self,
101        array: &V::Array,
102        parent: <Self::Parent as Matcher>::Match<'_>,
103        child_idx: usize,
104        ctx: &mut ExecutionCtx,
105    ) -> VortexResult<Option<ArrayRef>> {
106        assert_eq!(child_idx, 0);
107        if let Some(result) = precondition::<V>(array, &parent.range) {
108            return Ok(Some(result));
109        }
110        <V as SliceKernel>::slice(array, parent.range.clone(), ctx)
111    }
112}
vortex_array/arrays/slice/mod.rs

vortex_array/arrays/slice/
mod.rs
