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

vortex_array/arrays/slice/
vtable.rs
