vortex_array/arrays/primitive/array/
conversion.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Conversion methods and trait implementations of [`From`] and [`Into`] for [`PrimitiveArray`].
5
6use arrow_buffer::BooleanBufferBuilder;
7use vortex_buffer::{Buffer, BufferMut};
8use vortex_dtype::NativePType;
9use vortex_error::vortex_panic;
10
11use crate::arrays::PrimitiveArray;
12use crate::validity::Validity;
13use crate::vtable::ValidityHelper;
14use crate::{ArrayRef, IntoArray};
15
16impl PrimitiveArray {
17    /// Create a PrimitiveArray from an iterator of `T`.
18    /// NOTE: we cannot impl FromIterator trait since it conflicts with `FromIterator<T>`.
19    pub fn from_option_iter<T: NativePType, I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
20        let iter = iter.into_iter();
21        let mut values = BufferMut::with_capacity(iter.size_hint().0);
22        let mut validity = BooleanBufferBuilder::new(values.capacity());
23
24        for i in iter {
25            match i {
26                None => {
27                    validity.append(false);
28                    values.push(T::default());
29                }
30                Some(e) => {
31                    validity.append(true);
32                    values.push(e);
33                }
34            }
35        }
36        Self::new(values.freeze(), Validity::from(validity.finish()))
37    }
38
39    pub fn buffer<T: NativePType>(&self) -> Buffer<T> {
40        if T::PTYPE != self.ptype() {
41            vortex_panic!(
42                "Attempted to get buffer of type {} from array of type {}",
43                T::PTYPE,
44                self.ptype()
45            )
46        }
47        Buffer::from_byte_buffer(self.byte_buffer().clone())
48    }
49
50    pub fn into_buffer<T: NativePType>(self) -> Buffer<T> {
51        if T::PTYPE != self.ptype() {
52            vortex_panic!(
53                "Attempted to get buffer of type {} from array of type {}",
54                T::PTYPE,
55                self.ptype()
56            )
57        }
58        Buffer::from_byte_buffer(self.buffer)
59    }
60
61    /// Extract a mutable buffer from the PrimitiveArray. Attempts to do this with zero-copy
62    /// if the buffer is uniquely owned, otherwise will make a copy.
63    pub fn into_buffer_mut<T: NativePType>(self) -> BufferMut<T> {
64        if T::PTYPE != self.ptype() {
65            vortex_panic!(
66                "Attempted to get buffer_mut of type {} from array of type {}",
67                T::PTYPE,
68                self.ptype()
69            )
70        }
71        self.into_buffer()
72            .try_into_mut()
73            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
74    }
75
76    /// Try to extract a mutable buffer from the PrimitiveArray with zero copy.
77    #[allow(clippy::panic_in_result_fn)]
78    pub fn try_into_buffer_mut<T: NativePType>(self) -> Result<BufferMut<T>, PrimitiveArray> {
79        if T::PTYPE != self.ptype() {
80            vortex_panic!(
81                "Attempted to get buffer_mut of type {} from array of type {}",
82                T::PTYPE,
83                self.ptype()
84            )
85        }
86        let validity = self.validity().clone();
87        Buffer::<T>::from_byte_buffer(self.into_byte_buffer())
88            .try_into_mut()
89            .map_err(|buffer| PrimitiveArray::new(buffer, validity))
90    }
91}
92
93impl<T: NativePType> FromIterator<T> for PrimitiveArray {
94    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
95        let values = BufferMut::from_iter(iter);
96        PrimitiveArray::new(values, Validity::NonNullable)
97    }
98}
99
100impl<T: NativePType> IntoArray for Buffer<T> {
101    fn into_array(self) -> ArrayRef {
102        PrimitiveArray::new(self, Validity::NonNullable).into_array()
103    }
104}
105
106impl<T: NativePType> IntoArray for BufferMut<T> {
107    fn into_array(self) -> ArrayRef {
108        self.freeze().into_array()
109    }
110}
vortex_array/arrays/primitive/array/conversion.rs

vortex_array/arrays/primitive/array/
conversion.rs
