vortex_array/arrays/primitive/array/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use smallvec::smallvec;
8use vortex_buffer::Alignment;
9use vortex_buffer::Buffer;
10use vortex_buffer::BufferMut;
11use vortex_buffer::ByteBuffer;
12use vortex_buffer::ByteBufferMut;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_err;
16use vortex_error::vortex_panic;
17
18use crate::ArraySlots;
19use crate::ExecutionCtx;
20use crate::array::Array;
21use crate::array::ArrayParts;
22use crate::array::TypedArrayRef;
23use crate::arrays::BoolArray;
24use crate::arrays::Primitive;
25use crate::arrays::PrimitiveArray;
26use crate::dtype::DType;
27use crate::dtype::NativePType;
28use crate::dtype::Nullability;
29use crate::dtype::PType;
30use crate::match_each_native_ptype;
31use crate::validity::Validity;
32
33mod cast;
34mod conversion;
35mod patch;
36mod top_value;
37
38pub use patch::chunk_range;
39pub use patch::patch_chunk;
40
41use crate::ArrayRef;
42use crate::aggregate_fn::NumericalAggregateOpts;
43use crate::aggregate_fn::fns::min_max::min_max;
44use crate::array::child_to_validity;
45use crate::array::validity_to_child;
46use crate::array_slots;
47use crate::arrays::bool::BoolArrayExt;
48use crate::buffer::BufferHandle;
49use crate::builtins::ArrayBuiltins;
50
51#[array_slots(Primitive)]
52pub struct PrimitiveSlots {
53    /// The validity bitmap indicating which elements are non-null.
54    #[slot(0)]
55    pub validity: Option<ArrayRef>,
56}
57
58/// A primitive array that stores [native types][crate::dtype::NativePType] in a contiguous buffer
59/// of memory, along with an optional validity child.
60///
61/// This mirrors the Apache Arrow Primitive layout and can be converted into and out of one
62/// without allocations or copies.
63///
64/// The underlying buffer must be natively aligned to the primitive type they are representing.
65///
66/// Values are stored in their native representation with proper alignment.
67/// Null values still occupy space in the buffer but are marked invalid in the validity mask.
68///
69/// # Examples
70///
71/// ```
72/// # fn main() -> vortex_error::VortexResult<()> {
73/// use vortex_array::arrays::PrimitiveArray;
74/// use vortex_array::{VortexSessionExecute, array_session};
75///
76/// // Create from iterator using FromIterator impl
77/// let array: PrimitiveArray = [1i32, 2, 3, 4, 5].into_iter().collect();
78///
79/// // Slice the array
80/// let sliced = array.slice(1..3)?;
81///
82/// // Access individual values
83/// let mut ctx = array_session().create_execution_ctx();
84/// let value = sliced.execute_scalar(0, &mut ctx).unwrap();
85/// assert_eq!(value, 2i32.into());
86///
87/// # Ok(())
88/// # }
89/// ```
90#[derive(Clone, Debug)]
91pub struct PrimitiveData {
92    pub(super) ptype: PType,
93    pub(super) buffer: BufferHandle,
94}
95
96impl Display for PrimitiveData {
97    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
98        write!(f, "ptype: {}", self.ptype)
99    }
100}
101
102pub struct PrimitiveDataParts {
103    pub ptype: PType,
104    pub buffer: BufferHandle,
105    pub validity: Validity,
106}
107
108pub trait PrimitiveArrayExt: TypedArrayRef<Primitive> {
109    fn ptype(&self) -> PType {
110        match self.as_ref().dtype() {
111            DType::Primitive(ptype, _) => *ptype,
112            _ => unreachable!("PrimitiveArrayExt requires a primitive dtype"),
113        }
114    }
115
116    fn nullability(&self) -> Nullability {
117        match self.as_ref().dtype() {
118            DType::Primitive(_, nullability) => *nullability,
119            _ => unreachable!("PrimitiveArrayExt requires a primitive dtype"),
120        }
121    }
122
123    fn validity_child(&self) -> Option<&ArrayRef> {
124        self.as_ref().slots()[PrimitiveSlots::VALIDITY].as_ref()
125    }
126
127    fn validity(&self) -> Validity {
128        child_to_validity(
129            self.as_ref().slots()[PrimitiveSlots::VALIDITY].as_ref(),
130            self.nullability(),
131        )
132    }
133
134    fn buffer_handle(&self) -> &BufferHandle {
135        &self.buffer
136    }
137
138    fn reinterpret_cast(&self, ptype: PType) -> PrimitiveArray {
139        if self.ptype() == ptype {
140            return self.to_owned();
141        }
142
143        assert_eq!(
144            self.ptype().byte_width(),
145            ptype.byte_width(),
146            "can't reinterpret cast between integers of two different widths"
147        );
148
149        PrimitiveArray::from_buffer_handle(
150            self.buffer_handle().clone(),
151            ptype,
152            PrimitiveArrayExt::validity(self),
153        )
154    }
155
156    /// Narrow the array to the smallest possible integer type that can represent all values.
157    fn narrow(&self, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
158        if !self.ptype().is_int() {
159            return Ok(self.to_owned());
160        }
161
162        let Some(min_max) = min_max(self.as_ref(), ctx, NumericalAggregateOpts::default())? else {
163            return Ok(PrimitiveArray::new(
164                Buffer::<u8>::zeroed(self.len()),
165                PrimitiveArrayExt::validity(self),
166            ));
167        };
168
169        // If we can't cast to i64, then leave the array as its original type.
170        // It's too big to downcast anyway.
171        let Ok(min) = min_max
172            .min
173            .cast(&PType::I64.into())
174            .and_then(|s| i64::try_from(&s))
175        else {
176            return Ok(self.to_owned());
177        };
178        let Ok(max) = min_max
179            .max
180            .cast(&PType::I64.into())
181            .and_then(|s| i64::try_from(&s))
182        else {
183            return Ok(self.to_owned());
184        };
185
186        let nullability = self.as_ref().dtype().nullability();
187
188        if min < 0 || max < 0 {
189            // Signed
190            if min >= i8::MIN as i64 && max <= i8::MAX as i64 {
191                let result = self
192                    .as_ref()
193                    .cast(DType::Primitive(PType::I8, nullability))?
194                    .execute::<PrimitiveArray>(ctx)?;
195                return Ok(result);
196            }
197
198            if min >= i16::MIN as i64 && max <= i16::MAX as i64 {
199                let result = self
200                    .as_ref()
201                    .cast(DType::Primitive(PType::I16, nullability))?
202                    .execute::<PrimitiveArray>(ctx)?;
203                return Ok(result);
204            }
205
206            if min >= i32::MIN as i64 && max <= i32::MAX as i64 {
207                let result = self
208                    .as_ref()
209                    .cast(DType::Primitive(PType::I32, nullability))?
210                    .execute::<PrimitiveArray>(ctx)?;
211                return Ok(result);
212            }
213        } else {
214            // Unsigned
215            if max <= u8::MAX as i64 {
216                let result = self
217                    .as_ref()
218                    .cast(DType::Primitive(PType::U8, nullability))?
219                    .execute::<PrimitiveArray>(ctx)?;
220                return Ok(result);
221            }
222
223            if max <= u16::MAX as i64 {
224                let result = self
225                    .as_ref()
226                    .cast(DType::Primitive(PType::U16, nullability))?
227                    .execute::<PrimitiveArray>(ctx)?;
228                return Ok(result);
229            }
230
231            if max <= u32::MAX as i64 {
232                let result = self
233                    .as_ref()
234                    .cast(DType::Primitive(PType::U32, nullability))?
235                    .execute::<PrimitiveArray>(ctx)?;
236                return Ok(result);
237            }
238        }
239
240        Ok(self.to_owned())
241    }
242}
243impl<T: TypedArrayRef<Primitive>> PrimitiveArrayExt for T {}
244
245// TODO(connor): There are a lot of places where we could be using `new_unchecked` in the codebase.
246impl PrimitiveData {
247    /// Build the slots vector for this array.
248    pub(super) fn make_slots(validity: &Validity, len: usize) -> ArraySlots {
249        smallvec![validity_to_child(validity, len)]
250    }
251
252    /// Create a new array from a buffer handle.
253    ///
254    /// # Safety
255    ///
256    /// Should ensure that the provided BufferHandle points at sufficiently large region of aligned
257    /// memory to hold the `ptype` values.
258    pub unsafe fn new_unchecked_from_handle(
259        handle: BufferHandle,
260        ptype: PType,
261        _validity: Validity,
262    ) -> Self {
263        Self {
264            ptype,
265            buffer: handle,
266        }
267    }
268
269    /// Creates a new `PrimitiveArray`.
270    ///
271    /// # Panics
272    ///
273    /// Panics if the provided components do not satisfy the invariants documented
274    /// in `PrimitiveArray::new_unchecked`.
275    pub fn new<T: NativePType>(buffer: impl Into<Buffer<T>>, validity: Validity) -> Self {
276        let buffer = buffer.into();
277        Self::try_new(buffer, validity).vortex_expect("PrimitiveArray construction failed")
278    }
279
280    /// Constructs a new `PrimitiveArray`.
281    ///
282    /// See `PrimitiveArray::new_unchecked` for more information.
283    ///
284    /// # Errors
285    ///
286    /// Returns an error if the provided components do not satisfy the invariants documented in
287    /// `PrimitiveArray::new_unchecked`.
288    #[inline]
289    pub fn try_new<T: NativePType>(buffer: Buffer<T>, validity: Validity) -> VortexResult<Self> {
290        Self::validate(&buffer, &validity)?;
291
292        // SAFETY: validate ensures all invariants are met.
293        Ok(unsafe { Self::new_unchecked(buffer, validity) })
294    }
295
296    /// Creates a new `PrimitiveArray` without validation from these components:
297    ///
298    /// * `buffer` is a typed buffer containing the primitive values.
299    /// * `validity` holds the null values.
300    ///
301    /// # Safety
302    ///
303    /// The caller must ensure all of the following invariants are satisfied:
304    ///
305    /// ## Validity Requirements
306    ///
307    /// - If `validity` is [`Validity::Array`], its length must exactly equal `buffer.len()`.
308    #[inline]
309    pub unsafe fn new_unchecked<T: NativePType>(buffer: Buffer<T>, _validity: Validity) -> Self {
310        #[cfg(debug_assertions)]
311        Self::validate(&buffer, &_validity)
312            .vortex_expect("[Debug Assertion]: Invalid `PrimitiveArray` parameters");
313
314        Self {
315            ptype: T::PTYPE,
316            buffer: BufferHandle::new_host(buffer.into_byte_buffer()),
317        }
318    }
319
320    /// Validates the components that would be used to create a `PrimitiveArray`.
321    ///
322    /// This function checks all the invariants required by `PrimitiveArray::new_unchecked`.
323    #[inline]
324    pub fn validate<T: NativePType>(buffer: &Buffer<T>, validity: &Validity) -> VortexResult<()> {
325        if let Some(len) = validity.maybe_len()
326            && buffer.len() != len
327        {
328            return Err(vortex_err!(
329                InvalidArgument:
330                "Buffer and validity length mismatch: buffer={}, validity={}",
331                buffer.len(),
332                len
333            ));
334        }
335        Ok(())
336    }
337
338    pub fn empty<T: NativePType>(nullability: Nullability) -> Self {
339        Self::new(Buffer::<T>::empty(), nullability.into())
340    }
341}
342
343impl Array<Primitive> {
344    pub fn empty<T: NativePType>(nullability: Nullability) -> Self {
345        let dtype = DType::Primitive(T::PTYPE, nullability);
346        let len = 0;
347        let data = PrimitiveData::empty::<T>(nullability);
348        let slots = PrimitiveData::make_slots(&Validity::from(nullability), len);
349        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
350    }
351
352    /// Creates a new `PrimitiveArray`.
353    ///
354    /// # Panics
355    ///
356    /// Panics if the provided components do not satisfy the invariants.
357    pub fn new<T: NativePType>(buffer: impl Into<Buffer<T>>, validity: Validity) -> Self {
358        let buffer = buffer.into();
359        let dtype = DType::Primitive(T::PTYPE, validity.nullability());
360        let len = buffer.len();
361        let slots = PrimitiveData::make_slots(&validity, len);
362        let data = PrimitiveData::new(buffer, validity);
363        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
364    }
365
366    /// Constructs a new `PrimitiveArray`.
367    pub fn try_new<T: NativePType>(buffer: Buffer<T>, validity: Validity) -> VortexResult<Self> {
368        let dtype = DType::Primitive(T::PTYPE, validity.nullability());
369        let len = buffer.len();
370        let slots = PrimitiveData::make_slots(&validity, len);
371        let data = PrimitiveData::try_new(buffer, validity)?;
372        Ok(unsafe {
373            Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots))
374        })
375    }
376
377    /// Creates a new `PrimitiveArray` without validation.
378    ///
379    /// # Safety
380    ///
381    /// See [`PrimitiveData::new_unchecked`].
382    pub unsafe fn new_unchecked<T: NativePType>(buffer: Buffer<T>, validity: Validity) -> Self {
383        let dtype = DType::Primitive(T::PTYPE, validity.nullability());
384        let len = buffer.len();
385        let slots = PrimitiveData::make_slots(&validity, len);
386        let data = unsafe { PrimitiveData::new_unchecked(buffer, validity) };
387        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
388    }
389
390    /// Create a new array from a buffer handle.
391    ///
392    /// # Safety
393    ///
394    /// See [`PrimitiveData::new_unchecked_from_handle`].
395    pub unsafe fn new_unchecked_from_handle(
396        handle: BufferHandle,
397        ptype: PType,
398        validity: Validity,
399    ) -> Self {
400        let dtype = DType::Primitive(ptype, validity.nullability());
401        let len = handle.len() / ptype.byte_width();
402        let slots = PrimitiveData::make_slots(&validity, len);
403        let data = unsafe { PrimitiveData::new_unchecked_from_handle(handle, ptype, validity) };
404        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
405    }
406
407    /// Creates a new `PrimitiveArray` from a [`BufferHandle`].
408    pub fn from_buffer_handle(handle: BufferHandle, ptype: PType, validity: Validity) -> Self {
409        let dtype = DType::Primitive(ptype, validity.nullability());
410        let len = handle.len() / ptype.byte_width();
411        let slots = PrimitiveData::make_slots(&validity, len);
412        let data = PrimitiveData::from_buffer_handle(handle, ptype, validity);
413        Array::try_from_parts(ArrayParts::new(Primitive, dtype, len, data, slots))
414            .vortex_expect("PrimitiveData is always valid")
415    }
416
417    /// Creates a new `PrimitiveArray` from a [`ByteBuffer`].
418    pub fn from_byte_buffer(buffer: ByteBuffer, ptype: PType, validity: Validity) -> Self {
419        let dtype = DType::Primitive(ptype, validity.nullability());
420        let len = buffer.len() / ptype.byte_width();
421        let slots = PrimitiveData::make_slots(&validity, len);
422        let data = PrimitiveData::from_byte_buffer(buffer, ptype, validity);
423        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
424    }
425
426    /// Create a PrimitiveArray from a byte buffer containing only the valid elements.
427    pub fn from_values_byte_buffer(
428        valid_elems_buffer: ByteBuffer,
429        ptype: PType,
430        validity: Validity,
431        n_rows: usize,
432        ctx: &mut ExecutionCtx,
433    ) -> Self {
434        let dtype = DType::Primitive(ptype, validity.nullability());
435        let len = n_rows;
436        let slots = PrimitiveData::make_slots(&validity, len);
437        let data = PrimitiveData::from_values_byte_buffer(
438            valid_elems_buffer,
439            ptype,
440            validity,
441            n_rows,
442            ctx,
443        );
444        unsafe { Array::from_parts_unchecked(ArrayParts::new(Primitive, dtype, len, data, slots)) }
445    }
446
447    /// Validates the components that would be used to create a `PrimitiveArray`.
448    pub fn validate<T: NativePType>(buffer: &Buffer<T>, validity: &Validity) -> VortexResult<()> {
449        PrimitiveData::validate(buffer, validity)
450    }
451
452    pub fn into_data_parts(self) -> PrimitiveDataParts {
453        let validity = PrimitiveArrayExt::validity(&self);
454        let ptype = PrimitiveArrayExt::ptype(&self);
455        let data = self.into_data();
456        PrimitiveDataParts {
457            ptype,
458            buffer: data.buffer,
459            validity,
460        }
461    }
462}
463
464impl PrimitiveData {
465    pub fn len(&self) -> usize {
466        self.buffer.len() / self.ptype.byte_width()
467    }
468
469    /// Returns `true` if the array is empty.
470    pub fn is_empty(&self) -> bool {
471        self.buffer.is_empty()
472    }
473
474    pub fn ptype(&self) -> PType {
475        self.ptype
476    }
477
478    /// Get access to the buffer handle backing the array.
479    pub fn buffer_handle(&self) -> &BufferHandle {
480        &self.buffer
481    }
482
483    pub fn from_buffer_handle(handle: BufferHandle, ptype: PType, _validity: Validity) -> Self {
484        Self {
485            ptype,
486            buffer: handle,
487        }
488    }
489
490    pub fn from_byte_buffer(buffer: ByteBuffer, ptype: PType, validity: Validity) -> Self {
491        match_each_native_ptype!(ptype, |T| {
492            Self::new::<T>(Buffer::from_byte_buffer(buffer), validity)
493        })
494    }
495
496    /// Create a PrimitiveArray from a byte buffer containing only the valid elements.
497    pub fn from_values_byte_buffer(
498        valid_elems_buffer: ByteBuffer,
499        ptype: PType,
500        validity: Validity,
501        n_rows: usize,
502        ctx: &mut ExecutionCtx,
503    ) -> Self {
504        let byte_width = ptype.byte_width();
505        let alignment = Alignment::new(byte_width);
506        let buffer = match &validity {
507            Validity::AllValid | Validity::NonNullable => valid_elems_buffer.aligned(alignment),
508            Validity::AllInvalid => ByteBuffer::zeroed_aligned(n_rows * byte_width, alignment),
509            Validity::Array(is_valid) => {
510                let bool_array = is_valid
511                    .clone()
512                    .execute::<BoolArray>(ctx)
513                    .vortex_expect("must be a bool array");
514                let bool_buffer = bool_array.bit_buffer_view();
515                let mut bytes = ByteBufferMut::zeroed_aligned(n_rows * byte_width, alignment);
516                for (i, valid_i) in bool_buffer.set_indices().enumerate() {
517                    bytes[valid_i * byte_width..(valid_i + 1) * byte_width]
518                        .copy_from_slice(&valid_elems_buffer[i * byte_width..(i + 1) * byte_width])
519                }
520                bytes.freeze()
521            }
522        };
523
524        Self::from_byte_buffer(buffer, ptype, validity)
525    }
526
527    /// Get a buffer in host memory holding all the values.
528    ///
529    /// NOTE: some values may be nonsense if the validity buffer indicates that the value is null.
530    pub fn to_buffer<T: NativePType>(&self) -> Buffer<T> {
531        if T::PTYPE != self.ptype() {
532            vortex_panic!(
533                "Attempted to get buffer of type {} from array of type {}",
534                T::PTYPE,
535                self.ptype()
536            )
537        }
538        Buffer::from_byte_buffer(self.buffer_handle().to_host_sync())
539    }
540
541    /// Consume the array and get a host Buffer containing the data values.
542    pub fn into_buffer<T: NativePType>(self) -> Buffer<T> {
543        if T::PTYPE != self.ptype() {
544            vortex_panic!(
545                "Attempted to get buffer of type {} from array of type {}",
546                T::PTYPE,
547                self.ptype()
548            )
549        }
550        Buffer::from_byte_buffer(self.buffer.into_host_sync())
551    }
552
553    /// Extract a mutable buffer from the PrimitiveData. Attempts to do this with zero-copy
554    /// if the buffer is uniquely owned, otherwise will make a copy.
555    pub fn into_buffer_mut<T: NativePType>(self) -> BufferMut<T> {
556        self.try_into_buffer_mut()
557            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
558    }
559
560    /// Try to extract a mutable buffer from the PrimitiveData with zero copy.
561    ///
562    /// # Panic
563    /// If the buffer is not of type T this will panic
564    pub fn try_into_buffer_mut<T: NativePType>(self) -> Result<BufferMut<T>, Buffer<T>> {
565        if T::PTYPE != self.ptype() {
566            vortex_panic!(
567                "Attempted to get buffer_mut of type {} from array of type {}",
568                T::PTYPE,
569                self.ptype()
570            )
571        }
572        let buffer = Buffer::<T>::from_byte_buffer(self.buffer.into_host_sync());
573        buffer.try_into_mut()
574    }
575}
vortex_array/arrays/primitive/array/mod.rs

vortex_array/arrays/primitive/array/
mod.rs
