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

vortex_array/arrays/primitive/vtable/
mod.rs
