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_dtype::NativePType;
10use vortex_error::vortex_panic;
11
12use crate::ArrayRef;
13use crate::IntoArray;
14use crate::arrays::PrimitiveArray;
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        if T::PTYPE != self.ptype() {
58            vortex_panic!(
59                "Attempted to get buffer of type {} from array of type {}",
60                T::PTYPE,
61                self.ptype()
62            )
63        }
64
65        Buffer::from_byte_buffer(self.buffer.into_host_sync())
66    }
67
68    /// Extract a mutable buffer from the PrimitiveArray. Attempts to do this with zero-copy
69    /// if the buffer is uniquely owned, otherwise will make a copy.
70    pub fn into_buffer_mut<T: NativePType>(self) -> BufferMut<T> {
71        self.try_into_buffer_mut()
72            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
73    }
74
75    /// Try to extract a mutable buffer from the PrimitiveArray with zero copy.
76    pub fn try_into_buffer_mut<T: NativePType>(self) -> Result<BufferMut<T>, Buffer<T>> {
77        if T::PTYPE != self.ptype() {
78            vortex_panic!(
79                "Attempted to get buffer_mut of type {} from array of type {}",
80                T::PTYPE,
81                self.ptype()
82            )
83        }
84        let buffer = Buffer::<T>::from_byte_buffer(self.buffer.into_host_sync());
85        buffer.try_into_mut()
86    }
87}
88
89impl<T: NativePType> FromIterator<T> for PrimitiveArray {
90    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
91        let values = BufferMut::from_iter(iter);
92        PrimitiveArray::new(values, Validity::NonNullable)
93    }
94}
95
96impl<T: NativePType> IntoArray for Buffer<T> {
97    fn into_array(self) -> ArrayRef {
98        PrimitiveArray::new(self, Validity::NonNullable).into_array()
99    }
100}
101
102impl<T: NativePType> IntoArray for BufferMut<T> {
103    fn into_array(self) -> ArrayRef {
104        self.freeze().into_array()
105    }
106}