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
69fn precondition<V: VTable>(array: ArrayView<'_, V>, range: &Range<usize>) -> Option<ArrayRef> {
70    if range.start == 0 && range.end == array.len() {
71        return Some(array.array().clone());
72    };
73    if range.start == range.end {
74        return Some(Canonical::empty(array.dtype()).into_array());
75    }
76    None
77}
78
79/// Adaptor that wraps a [`SliceReduce`] impl as an [`ArrayParentReduceRule`].
80#[derive(Default, Debug)]
81pub struct SliceReduceAdaptor<V>(pub V);
82
83impl<V> ArrayParentReduceRule<V> for SliceReduceAdaptor<V>
84where
85    V: SliceReduce,
86{
87    type Parent = Slice;
88
89    fn reduce_parent(
90        &self,
91        array: ArrayView<'_, V>,
92        parent: <Self::Parent as Matcher>::Match<'_>,
93        child_idx: usize,
94    ) -> VortexResult<Option<ArrayRef>> {
95        assert_eq!(child_idx, 0);
96        if let Some(result) = precondition::<V>(array, &parent.range) {
97            return Ok(Some(result));
98        }
99        <V as SliceReduce>::slice(array, parent.range.clone())
100    }
101}
102
103/// Adaptor that wraps a [`SliceKernel`] impl as an [`ExecuteParentKernel`].
104#[derive(Default, Debug)]
105pub struct SliceExecuteAdaptor<V>(pub V);
106
107impl<V> ExecuteParentKernel<V> for SliceExecuteAdaptor<V>
108where
109    V: SliceKernel,
110{
111    type Parent = Slice;
112
113    fn execute_parent(
114        &self,
115        array: ArrayView<'_, V>,
116        parent: <Self::Parent as Matcher>::Match<'_>,
117        child_idx: usize,
118        ctx: &mut ExecutionCtx,
119    ) -> VortexResult<Option<ArrayRef>> {
120        assert_eq!(child_idx, 0);
121        if let Some(result) = precondition::<V>(array, &parent.range) {
122            return Ok(Some(result));
123        }
124        <V as SliceKernel>::slice(array, parent.range.clone(), ctx)
125    }
126}
vortex_array/arrays/slice/mod.rs

vortex_array/arrays/slice/
mod.rs
