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::ArrayEq;
19use crate::ArrayHash;
20use crate::ArrayRef;
21use crate::Precision;
22use crate::array::Array;
23use crate::array::ArrayId;
24use crate::array::ArrayView;
25use crate::array::OperationsVTable;
26use crate::array::VTable;
27use crate::array::ValidityVTable;
28use crate::arrays::slice::SliceArrayExt;
29use crate::arrays::slice::array::CHILD_SLOT;
30use crate::arrays::slice::array::SLOT_NAMES;
31use crate::arrays::slice::array::SliceData;
32use crate::arrays::slice::rules::PARENT_RULES;
33use crate::buffer::BufferHandle;
34use crate::dtype::DType;
35use crate::executor::ExecutionCtx;
36use crate::executor::ExecutionResult;
37use crate::require_child;
38use crate::scalar::Scalar;
39use crate::serde::ArrayChildren;
40use crate::validity::Validity;
41
42/// A [`Slice`]-encoded Vortex array.
43pub type SliceArray = Array<Slice>;
44
45#[derive(Clone, Debug)]
46pub struct Slice;
47
48impl Slice {
49    pub const ID: ArrayId = ArrayId::new_ref("vortex.slice");
50}
51
52impl ArrayHash for SliceData {
53    fn array_hash<H: Hasher>(&self, state: &mut H, _precision: Precision) {
54        self.range.start.hash(state);
55        self.range.end.hash(state);
56    }
57}
58
59impl ArrayEq for SliceData {
60    fn array_eq(&self, other: &Self, _precision: Precision) -> bool {
61        self.range == other.range
62    }
63}
64
65impl VTable for Slice {
66    type ArrayData = SliceData;
67    type OperationsVTable = Self;
68    type ValidityVTable = Self;
69
70    fn id(&self) -> ArrayId {
71        Slice::ID
72    }
73
74    fn validate(
75        &self,
76        data: &Self::ArrayData,
77        dtype: &DType,
78        len: usize,
79        slots: &[Option<ArrayRef>],
80    ) -> VortexResult<()> {
81        vortex_ensure!(
82            slots[CHILD_SLOT].is_some(),
83            "SliceArray child slot must be present"
84        );
85        let child = slots[CHILD_SLOT]
86            .as_ref()
87            .vortex_expect("validated child slot");
88        vortex_ensure!(
89            child.dtype() == dtype,
90            "SliceArray dtype {} does not match outer dtype {}",
91            child.dtype(),
92            dtype
93        );
94        vortex_ensure!(
95            data.len() == len,
96            "SliceArray length {} does not match outer length {}",
97            data.len(),
98            len
99        );
100        vortex_ensure!(
101            data.range.end <= child.len(),
102            "SliceArray range {:?} exceeds child length {}",
103            data.range,
104            child.len()
105        );
106        Ok(())
107    }
108
109    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
110        0
111    }
112
113    fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
114        vortex_panic!("SliceArray has no buffers")
115    }
116
117    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
118        None
119    }
120
121    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
122        SLOT_NAMES[idx].to_string()
123    }
124
125    fn serialize(
126        _array: ArrayView<'_, Self>,
127        _session: &VortexSession,
128    ) -> VortexResult<Option<Vec<u8>>> {
129        // TODO(joe): make this configurable
130        vortex_bail!("Slice array is not serializable")
131    }
132
133    fn deserialize(
134        &self,
135        _dtype: &DType,
136        _len: usize,
137        _metadata: &[u8],
138
139        _buffers: &[BufferHandle],
140        _children: &dyn ArrayChildren,
141        _session: &VortexSession,
142    ) -> VortexResult<crate::array::ArrayParts<Self>> {
143        vortex_bail!("Slice array is not serializable")
144    }
145
146    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
147        let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical);
148
149        debug_assert!(array.child().is_canonical());
150        // TODO(ngates): we should inline canonical slice logic here.
151        array
152            .child()
153            .slice(array.range.clone())
154            .map(ExecutionResult::done)
155    }
156
157    fn reduce_parent(
158        array: ArrayView<'_, Self>,
159        parent: &ArrayRef,
160        child_idx: usize,
161    ) -> VortexResult<Option<ArrayRef>> {
162        PARENT_RULES.evaluate(array, parent, child_idx)
163    }
164}
165impl OperationsVTable<Slice> for Slice {
166    fn scalar_at(
167        array: ArrayView<'_, Slice>,
168        index: usize,
169        _ctx: &mut ExecutionCtx,
170    ) -> VortexResult<Scalar> {
171        array.child().scalar_at(array.range.start + index)
172    }
173}
174
175impl ValidityVTable<Slice> for Slice {
176    fn validity(array: ArrayView<'_, Slice>) -> VortexResult<Validity> {
177        array.child().validity()?.slice(array.range.clone())
178    }
179}
180
181pub struct SliceMetadata(pub(super) Range<usize>);
182
183impl Debug for SliceMetadata {
184    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
185        write!(f, "{}..{}", self.0.start, self.0.end)
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use vortex_error::VortexResult;
192
193    use crate::IntoArray;
194    use crate::arrays::PrimitiveArray;
195    use crate::arrays::SliceArray;
196    use crate::assert_arrays_eq;
197
198    #[test]
199    fn test_slice_slice() -> VortexResult<()> {
200        // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base)
201        let arr = PrimitiveArray::from_iter(0i32..10).into_array();
202        let inner_slice = SliceArray::new(arr, 2..8).into_array();
203        let slice = inner_slice.slice(1..4)?;
204
205        assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5]));
206
207        Ok(())
208    }
209}
vortex_array/arrays/slice/vtable.rs

vortex_array/arrays/slice/
vtable.rs
