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::SliceArraySlotsExt;
19pub use array::SliceData;
20pub use array::SliceDataParts;
21pub use array::SliceSlots;
22pub use array::SliceSlotsView;
23use vortex_error::VortexResult;
24pub use vtable::*;
25
26use crate::ArrayRef;
27use crate::Canonical;
28use crate::ExecutionCtx;
29use crate::IntoArray;
30use crate::array::ArrayView;
31use crate::array::VTable;
32use crate::kernel::ExecuteParentKernel;
33use crate::matcher::Matcher;
34use crate::optimizer::rules::ArrayParentReduceRule;
35
36pub trait SliceReduce: VTable {
37    /// Slice an array with the provided range without reading buffers.
38    ///
39    /// This trait is for slice implementations that can operate purely on array metadata and
40    /// structure without needing to read or execute on the underlying buffers. Implementations
41    /// should return `None` if slicing requires buffer access.
42    ///
43    /// # Preconditions
44    ///
45    /// The range is guaranteed to be within bounds of the array (i.e., `range.end <= array.len()`).
46    ///
47    /// Additionally, the range is guaranteed to be non-empty (i.e., `range.start < range.end`).
48    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>>;
49}
50
51pub trait SliceKernel: VTable {
52    /// Slice an array with the provided range, potentially reading buffers.
53    ///
54    /// Unlike [`SliceReduce`], this trait is for slice implementations that may need to read
55    /// and execute on the underlying buffers to produce the sliced result.
56    ///
57    /// # Preconditions
58    ///
59    /// The range is guaranteed to be within bounds of the array (i.e., `range.end <= array.len()`).
60    ///
61    /// Additionally, the range is guaranteed to be non-empty (i.e., `range.start < range.end`).
62    fn slice(
63        array: ArrayView<'_, Self>,
64        range: Range<usize>,
65        ctx: &mut ExecutionCtx,
66    ) -> VortexResult<Option<ArrayRef>>;
67}
68
69/// Short-circuits slice for the ranges that need no encoding-specific work.
70///
71/// Returns `Some(result)` when the answer is already known, or `None` when the slice must proceed
72/// normally.
73fn short_circuit<V: VTable>(array: ArrayView<'_, V>, range: &Range<usize>) -> Option<ArrayRef> {
74    if range.start == 0 && range.end == array.len() {
75        return Some(array.array().clone());
76    }
77    if range.start == range.end {
78        return Some(Canonical::empty(array.dtype()).into_array());
79    }
80    None
81}
82
83/// Adaptor that wraps a [`SliceReduce`] impl as an [`ArrayParentReduceRule`].
84#[derive(Default, Debug)]
85pub struct SliceReduceAdaptor<V>(pub V);
86
87impl<V> ArrayParentReduceRule<V> for SliceReduceAdaptor<V>
88where
89    V: SliceReduce,
90{
91    type Parent = Slice;
92
93    fn reduce_parent(
94        &self,
95        array: ArrayView<'_, V>,
96        parent: <Self::Parent as Matcher>::Match<'_>,
97        child_idx: usize,
98    ) -> VortexResult<Option<ArrayRef>> {
99        assert_eq!(child_idx, 0);
100        if let Some(result) = short_circuit::<V>(array, &parent.range) {
101            return Ok(Some(result));
102        }
103        <V as SliceReduce>::slice(array, parent.range.clone())
104    }
105}
106
107/// Adaptor that wraps a [`SliceKernel`] impl as an [`ExecuteParentKernel`].
108#[derive(Default, Debug)]
109pub struct SliceExecuteAdaptor<V>(pub V);
110
111impl<V> ExecuteParentKernel<V> for SliceExecuteAdaptor<V>
112where
113    V: SliceKernel,
114{
115    type Parent = Slice;
116
117    fn execute_parent(
118        &self,
119        array: ArrayView<'_, V>,
120        parent: <Self::Parent as Matcher>::Match<'_>,
121        child_idx: usize,
122        ctx: &mut ExecutionCtx,
123    ) -> VortexResult<Option<ArrayRef>> {
124        assert_eq!(child_idx, 0);
125        if let Some(result) = short_circuit::<V>(array, &parent.range) {
126            return Ok(Some(result));
127        }
128        <V as SliceKernel>::slice(array, parent.range.clone(), ctx)
129    }
130}
vortex_array/arrays/slice/mod.rs

vortex_array/arrays/slice/
mod.rs
