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

vortex_array/arrays/slice/
mod.rs
