vortex_array/arrays/primitive/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_error::vortex_ensure;
7use vortex_error::vortex_ensure_eq;
8
9use crate::ArrayParts;
10use crate::ArrayRef;
11use crate::ExecutionCtx;
12use crate::ExecutionResult;
13use crate::array::Array;
14use crate::array::ArrayView;
15use crate::array::VTable;
16use crate::arrays::fixed_width::vtable as fixed_width;
17use crate::arrays::primitive::PrimitiveData;
18use crate::buffer::BufferHandle;
19use crate::builders::ArrayBuilder;
20use crate::builders::PrimitiveBuilder;
21use crate::dtype::DType;
22use crate::dtype::PType;
23use crate::match_each_native_ptype;
24use crate::serde::ArrayChildren;
25mod kernel;
26mod operations;
27mod validity;
28
29use std::hash::Hasher;
30
31use vortex_buffer::Alignment;
32use vortex_session::VortexSession;
33use vortex_session::registry::CachedId;
34
35use crate::EqMode;
36use crate::array::ArrayId;
37use crate::arrays::primitive::array::PrimitiveSlots;
38use crate::arrays::primitive::compute::rules::RULES;
39use crate::hash::ArrayEq;
40use crate::hash::ArrayHash;
41
42/// A [`Primitive`]-encoded Vortex array.
43pub type PrimitiveArray = Array<Primitive>;
44
45pub(crate) fn initialize(session: &VortexSession) {
46    kernel::initialize(session);
47}
48
49impl ArrayHash for PrimitiveData {
50    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
51        self.buffer.array_hash(state, accuracy);
52    }
53}
54
55impl ArrayEq for PrimitiveData {
56    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
57        self.buffer.array_eq(&other.buffer, accuracy)
58    }
59}
60
61impl VTable for Primitive {
62    type TypedArrayData = PrimitiveData;
63
64    type OperationsVTable = Self;
65    type ValidityVTable = Self;
66
67    fn id(&self) -> ArrayId {
68        static ID: CachedId = CachedId::new("vortex.primitive");
69        *ID
70    }
71
72    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
73        1
74    }
75
76    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
77        fixed_width::buffer("PrimitiveArray", array.buffer_handle(), idx)
78    }
79
80    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
81        fixed_width::buffer_name(idx)
82    }
83
84    fn with_buffers(
85        &self,
86        array: ArrayView<'_, Self>,
87        buffers: &[BufferHandle],
88    ) -> VortexResult<ArrayParts<Self>> {
89        let mut data = array.data().clone();
90        data.buffer = fixed_width::single_buffer(buffers)?;
91        Ok(ArrayParts::new(
92            self.clone(),
93            array.dtype().clone(),
94            array.len(),
95            data,
96            array.slots().iter().cloned().collect(),
97        ))
98    }
99
100    fn serialize(
101        _array: ArrayView<'_, Self>,
102        _session: &VortexSession,
103    ) -> VortexResult<Option<Vec<u8>>> {
104        Ok(Some(vec![]))
105    }
106
107    fn validate(
108        &self,
109        data: &PrimitiveData,
110        dtype: &DType,
111        len: usize,
112        slots: &[Option<ArrayRef>],
113    ) -> VortexResult<()> {
114        let DType::Primitive(_, nullability) = dtype else {
115            vortex_bail!("Expected primitive dtype, got {dtype:?}");
116        };
117        vortex_ensure_eq!(
118            data.len(),
119            len,
120            "PrimitiveArray length does not match outer length",
121        );
122        let validity =
123            crate::array::child_to_validity(slots[PrimitiveSlots::VALIDITY].as_ref(), *nullability);
124        if let Some(validity_len) = validity.maybe_len() {
125            vortex_ensure_eq!(
126                validity_len,
127                len,
128                "PrimitiveArray validity len does not match outer length",
129            );
130        }
131
132        Ok(())
133    }
134
135    fn deserialize(
136        &self,
137        dtype: &DType,
138        len: usize,
139        metadata: &[u8],
140
141        buffers: &[BufferHandle],
142        children: &dyn ArrayChildren,
143        _session: &VortexSession,
144    ) -> VortexResult<ArrayParts<Self>> {
145        if !metadata.is_empty() {
146            vortex_bail!(
147                "PrimitiveArray expects empty metadata, got {} bytes",
148                metadata.len()
149            );
150        }
151        let buffer = fixed_width::single_buffer(buffers)?;
152
153        let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?;
154
155        let ptype = PType::try_from(dtype)?;
156
157        vortex_ensure!(
158            buffer.is_aligned_to(Alignment::new(ptype.byte_width())),
159            "Misaligned buffer cannot be used to build PrimitiveArray of {ptype}"
160        );
161
162        if buffer.len() != ptype.byte_width() * len {
163            vortex_bail!(
164                "Buffer length {} does not match expected length {} for {}, {}",
165                buffer.len(),
166                ptype.byte_width() * len,
167                ptype.byte_width(),
168                len,
169            );
170        }
171
172        // SAFETY: the buffer length and alignment are checked above.
173        let slots = PrimitiveData::make_slots(&validity, len);
174        let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) };
175        Ok(ArrayParts::new(
176            self.clone(),
177            dtype.clone(),
178            len,
179            data,
180            slots,
181        ))
182    }
183
184    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
185        PrimitiveSlots::NAMES[idx].to_string()
186    }
187
188    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
189        Ok(ExecutionResult::done(array))
190    }
191
192    fn append_to_builder(
193        array: ArrayView<'_, Self>,
194        builder: &mut dyn ArrayBuilder,
195        ctx: &mut ExecutionCtx,
196    ) -> VortexResult<()> {
197        match_each_native_ptype!(array.ptype(), |P| {
198            if let Some(builder) = builder.as_any_mut().downcast_mut::<PrimitiveBuilder<P>>() {
199                return builder.append_primitive_array(&array.into_owned(), ctx);
200            }
201        });
202
203        vortex_bail!("append_to_builder for Primitive requires a matching PrimitiveBuilder");
204    }
205
206    fn reduce_parent(
207        array: ArrayView<'_, Self>,
208        parent: &ArrayRef,
209        child_idx: usize,
210    ) -> VortexResult<Option<ArrayRef>> {
211        RULES.evaluate(array, parent, child_idx)
212    }
213}
214
215#[derive(Clone, Debug)]
216pub struct Primitive;
217
218#[cfg(test)]
219mod tests {
220    use vortex_buffer::ByteBufferMut;
221    use vortex_buffer::buffer;
222    use vortex_error::VortexResult;
223    use vortex_session::registry::ReadContext;
224
225    use crate::ArrayContext;
226    use crate::IntoArray;
227    use crate::VortexSessionExecute;
228    use crate::array_session;
229    use crate::arrays::PrimitiveArray;
230    use crate::assert_arrays_eq;
231    use crate::buffer::BufferHandle;
232    use crate::serde::SerializeOptions;
233    use crate::serde::SerializedArray;
234    use crate::validity::Validity;
235
236    #[test]
237    fn test_nullable_primitive_serde_roundtrip() {
238        let session = array_session();
239        let mut ctx = session.create_execution_ctx();
240        let array = PrimitiveArray::new(
241            buffer![1i32, 2, 3, 4],
242            Validity::from_iter([true, false, true, false]),
243        );
244        let dtype = array.dtype().clone();
245        let len = array.len();
246
247        let array_ctx = ArrayContext::empty();
248        let serialized = array
249            .clone()
250            .into_array()
251            .serialize(&array_ctx, &session, &SerializeOptions::default())
252            .unwrap();
253
254        let mut concat = ByteBufferMut::empty();
255        for buf in serialized {
256            concat.extend_from_slice(buf.as_ref());
257        }
258        let parts = SerializedArray::try_from(concat.freeze()).unwrap();
259        let decoded = parts
260            .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)
261            .unwrap();
262
263        assert_arrays_eq!(decoded, array, &mut ctx);
264    }
265
266    #[test]
267    fn test_with_buffers_replaces_primitive_buffer_with_equivalent_contents() -> VortexResult<()> {
268        let session = array_session();
269        let mut ctx = session.create_execution_ctx();
270
271        let array = PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array();
272        let replacement = BufferHandle::new_host(buffer![1i32, 2, 3, 4].into_byte_buffer());
273        // SAFETY: the replacement buffer contains the same logical values as the original array;
274        // only the buffer handle changes.
275        let rewritten = unsafe { array.with_buffers([replacement]) }?;
276        let expected = PrimitiveArray::from_iter([1i32, 2, 3, 4]);
277
278        assert_arrays_eq!(rewritten, expected, &mut ctx);
279        Ok(())
280    }
281
282    #[test]
283    fn test_with_buffers_rejects_length_change() {
284        let array = PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array();
285        let replacement = BufferHandle::new_host(buffer![10i32, 20, 30].into_byte_buffer());
286
287        // SAFETY: this call is expected to fail the checked buffer length invariant before any
288        // rewritten array is returned or observed.
289        assert!(unsafe { array.with_buffers([replacement]) }.is_err());
290    }
291}
vortex_array/arrays/primitive/vtable/mod.rs

vortex_array/arrays/primitive/vtable/
mod.rs
