vortex_array/arrays/slice/
mod.rs

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

vortex_array/arrays/slice/
mod.rs
