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

vortex_array/arrays/slice/
vtable.rs
