vortex_array/arrays/slice/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Formatter;
6use std::hash::Hash;
7use std::hash::Hasher;
8use std::ops::Range;
9use std::sync::Arc;
10
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_ensure;
15use vortex_error::vortex_panic;
16use vortex_session::VortexSession;
17
18use crate::AnyCanonical;
19use crate::ArrayEq;
20use crate::ArrayHash;
21use crate::ArrayRef;
22use crate::Canonical;
23use crate::DynArray;
24use crate::Precision;
25use crate::arrays::slice::array::SliceArray;
26use crate::arrays::slice::rules::PARENT_RULES;
27use crate::buffer::BufferHandle;
28use crate::dtype::DType;
29use crate::executor::ExecutionCtx;
30use crate::executor::ExecutionResult;
31use crate::scalar::Scalar;
32use crate::serde::ArrayChildren;
33use crate::stats::StatsSetRef;
34use crate::validity::Validity;
35use crate::vtable;
36use crate::vtable::ArrayId;
37use crate::vtable::OperationsVTable;
38use crate::vtable::VTable;
39use crate::vtable::ValidityVTable;
40
41vtable!(Slice);
42
43#[derive(Clone, Debug)]
44pub struct Slice;
45
46impl Slice {
47    pub const ID: ArrayId = ArrayId::new_ref("vortex.slice");
48}
49
50impl VTable for Slice {
51    type Array = SliceArray;
52    type Metadata = SliceMetadata;
53    type OperationsVTable = Self;
54    type ValidityVTable = Self;
55    fn vtable(_array: &Self::Array) -> &Self {
56        &Slice
57    }
58
59    fn id(&self) -> ArrayId {
60        Slice::ID
61    }
62
63    fn len(array: &SliceArray) -> usize {
64        array.range.len()
65    }
66
67    fn dtype(array: &SliceArray) -> &DType {
68        array.child.dtype()
69    }
70
71    fn stats(array: &SliceArray) -> StatsSetRef<'_> {
72        array.stats.to_ref(array.as_ref())
73    }
74
75    fn array_hash<H: Hasher>(array: &SliceArray, state: &mut H, precision: Precision) {
76        array.child.array_hash(state, precision);
77        array.range.start.hash(state);
78        array.range.end.hash(state);
79    }
80
81    fn array_eq(array: &SliceArray, other: &SliceArray, precision: Precision) -> bool {
82        array.child.array_eq(&other.child, precision) && array.range == other.range
83    }
84
85    fn nbuffers(_array: &Self::Array) -> usize {
86        0
87    }
88
89    fn buffer(_array: &Self::Array, _idx: usize) -> BufferHandle {
90        vortex_panic!("SliceArray has no buffers")
91    }
92
93    fn buffer_name(_array: &Self::Array, _idx: usize) -> Option<String> {
94        None
95    }
96
97    fn nchildren(_array: &Self::Array) -> usize {
98        1
99    }
100
101    fn child(array: &Self::Array, idx: usize) -> ArrayRef {
102        match idx {
103            0 => array.child.clone(),
104            _ => vortex_panic!("SliceArray child index {idx} out of bounds"),
105        }
106    }
107
108    fn child_name(_array: &Self::Array, idx: usize) -> String {
109        match idx {
110            0 => "child".to_string(),
111            _ => vortex_panic!("SliceArray child_name index {idx} out of bounds"),
112        }
113    }
114
115    fn metadata(array: &Self::Array) -> VortexResult<Self::Metadata> {
116        Ok(SliceMetadata(array.range.clone()))
117    }
118
119    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
120        // TODO(joe): make this configurable
121        vortex_bail!("Slice array is not serializable")
122    }
123
124    fn deserialize(
125        _bytes: &[u8],
126        _dtype: &DType,
127        _len: usize,
128        _buffers: &[BufferHandle],
129        _session: &VortexSession,
130    ) -> VortexResult<Self::Metadata> {
131        vortex_bail!("Slice array is not serializable")
132    }
133
134    fn build(
135        dtype: &DType,
136        len: usize,
137        metadata: &SliceMetadata,
138        _buffers: &[BufferHandle],
139        children: &dyn ArrayChildren,
140    ) -> VortexResult<Self::Array> {
141        assert_eq!(len, metadata.0.len());
142        let child = children.get(0, dtype, metadata.0.end)?;
143        Ok(SliceArray {
144            child,
145            range: metadata.0.clone(),
146            stats: Default::default(),
147        })
148    }
149
150    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
151        vortex_ensure!(
152            children.len() == 1,
153            "SliceArray expects exactly 1 child, got {}",
154            children.len()
155        );
156        array.child = children
157            .into_iter()
158            .next()
159            .vortex_expect("children length already validated");
160        Ok(())
161    }
162
163    fn execute(array: Arc<Self::Array>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
164        // Execute the child to get canonical form, then slice it
165        let Some(canonical) = array.child.as_opt::<AnyCanonical>() else {
166            // If the child is not canonical, recurse.
167            return array
168                .child
169                .clone()
170                .execute::<ArrayRef>(ctx)?
171                .slice(array.slice_range().clone())
172                .map(ExecutionResult::done);
173        };
174
175        // TODO(ngates): we should inline canonical slice logic here.
176        Canonical::from(canonical)
177            .as_ref()
178            .slice(array.range.clone())
179            .map(ExecutionResult::done)
180    }
181
182    fn reduce_parent(
183        array: &Self::Array,
184        parent: &ArrayRef,
185        child_idx: usize,
186    ) -> VortexResult<Option<ArrayRef>> {
187        PARENT_RULES.evaluate(array, parent, child_idx)
188    }
189}
190impl OperationsVTable<Slice> for Slice {
191    fn scalar_at(array: &SliceArray, index: usize) -> VortexResult<Scalar> {
192        array.child.scalar_at(array.range.start + index)
193    }
194}
195
196impl ValidityVTable<Slice> for Slice {
197    fn validity(array: &SliceArray) -> VortexResult<Validity> {
198        array.child.validity()?.slice(array.range.clone())
199    }
200}
201
202pub struct SliceMetadata(pub(super) Range<usize>);
203
204impl Debug for SliceMetadata {
205    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
206        write!(f, "{}..{}", self.0.start, self.0.end)
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use vortex_error::VortexResult;
213
214    use crate::DynArray;
215    use crate::IntoArray;
216    use crate::arrays::PrimitiveArray;
217    use crate::arrays::SliceArray;
218    use crate::assert_arrays_eq;
219
220    #[test]
221    fn test_slice_slice() -> VortexResult<()> {
222        // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base)
223        let arr = PrimitiveArray::from_iter(0i32..10).into_array();
224        let inner_slice = SliceArray::new(arr, 2..8).into_array();
225        let slice = inner_slice.slice(1..4)?;
226
227        assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5]));
228
229        Ok(())
230    }
231}
vortex_array/arrays/slice/vtable.rs

vortex_array/arrays/slice/
vtable.rs
