Skip to main content

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 vortex_buffer::BitBufferMut;
7use vortex_buffer::Buffer;
8use vortex_buffer::BufferMut;
9use vortex_error::vortex_panic;
10
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::arrays::PrimitiveArray;
14use crate::dtype::NativePType;
15use crate::validity::Validity;
16
17impl PrimitiveArray {
18    /// Create a PrimitiveArray from an iterator of `T`.
19    /// NOTE: we cannot impl FromIterator trait since it conflicts with `FromIterator<T>`.
20    pub fn from_option_iter<T: NativePType, I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
21        let iter = iter.into_iter();
22        let mut values = BufferMut::with_capacity(iter.size_hint().0);
23        let mut validity = BitBufferMut::with_capacity(values.capacity());
24
25        for i in iter {
26            match i {
27                None => {
28                    validity.append(false);
29                    values.push(T::default());
30                }
31                Some(e) => {
32                    validity.append(true);
33                    values.push(e);
34                }
35            }
36        }
37        Self::new(values.freeze(), Validity::from(validity.freeze()))
38    }
39
40    /// Get a buffer in host memory holding all the values.
41    ///
42    /// NOTE: some values may be nonsense if the validity buffer indicates that the value is null.
43    pub fn to_buffer<T: NativePType>(&self) -> Buffer<T> {
44        if T::PTYPE != self.ptype() {
45            vortex_panic!(
46                "Attempted to get buffer of type {} from array of type {}",
47                T::PTYPE,
48                self.ptype()
49            )
50        }
51
52        Buffer::from_byte_buffer(self.buffer_handle().to_host_sync())
53    }
54
55    /// Consume the array and get a host Buffer containing the data values.
56    pub fn into_buffer<T: NativePType>(self) -> Buffer<T> {
57        self.into_data().into_buffer()
58    }
59
60    /// Extract a mutable buffer from the PrimitiveArray. Attempts to do this with zero-copy
61    /// if the buffer is uniquely owned, otherwise will make a copy.
62    pub fn into_buffer_mut<T: NativePType>(self) -> BufferMut<T> {
63        self.into_data().into_buffer_mut()
64    }
65
66    /// Try to extract a mutable buffer from the PrimitiveArray with zero copy.
67    pub fn try_into_buffer_mut<T: NativePType>(self) -> Result<BufferMut<T>, Buffer<T>> {
68        self.into_data().try_into_buffer_mut()
69    }
70}
71
72impl<T: NativePType> FromIterator<T> for PrimitiveArray {
73    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
74        let values = BufferMut::from_iter(iter);
75        PrimitiveArray::new(values, Validity::NonNullable)
76    }
77}
78
79impl<T: NativePType> IntoArray for Buffer<T> {
80    fn into_array(self) -> ArrayRef {
81        PrimitiveArray::new(self, Validity::NonNullable).into_array()
82    }
83}
84
85impl<T: NativePType> IntoArray for BufferMut<T> {
86    fn into_array(self) -> ArrayRef {
87        self.freeze().into_array()
88    }
89}