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

vortex_array/arrays/slice/
vtable.rs
