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

vortex_array/arrays/slice/
vtable.rs
