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, Buffer, BufferMut};
7use vortex_dtype::{NativePType, Nullability};
8use vortex_error::{VortexResult, vortex_ensure, vortex_panic};
9use vortex_vector::primitive::PrimitiveVector;
10use vortex_vector::{VectorOps, match_each_pvector};
11
12use crate::arrays::PrimitiveArray;
13use crate::validity::Validity;
14use crate::vtable::ValidityHelper;
15use crate::{ArrayRef, IntoArray};
16
17impl PrimitiveArray {
18    /// Attempts to create a `PrimitiveArray` from a [`PrimitiveVector`] given a [`Nullability`].
19    ///
20    /// # Errors
21    ///
22    /// Returns an error if the nullability is [`NonNullable`](Nullability::NonNullable) and there
23    /// are nulls present in the vector.
24    pub fn try_from_vector(
25        primitive_vector: PrimitiveVector,
26        nullability: Nullability,
27    ) -> VortexResult<Self> {
28        // If we want to create a non-nullable array, then the vector should not have any nulls.
29        vortex_ensure!(
30            nullability.is_nullable() || primitive_vector.validity().all_true(),
31            "tried to create a non-nullable `PrimitiveArray` from a `PrimitiveVector` that had nulls"
32        );
33
34        match_each_pvector!(primitive_vector, |v| {
35            let (buffer, mask) = v.into_parts();
36            debug_assert_eq!(buffer.len(), mask.len());
37
38            let validity = Validity::from_mask(mask, nullability);
39
40            // SAFETY: Since the buffer and the mask came from a valid vector, we know that the
41            // length of the buffer and the validity are the same.
42            Ok(unsafe { Self::new_unchecked(buffer, validity) })
43        })
44    }
45
46    /// Create a PrimitiveArray from an iterator of `T`.
47    /// NOTE: we cannot impl FromIterator trait since it conflicts with `FromIterator<T>`.
48    pub fn from_option_iter<T: NativePType, I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
49        let iter = iter.into_iter();
50        let mut values = BufferMut::with_capacity(iter.size_hint().0);
51        let mut validity = BitBufferMut::with_capacity(values.capacity());
52
53        for i in iter {
54            match i {
55                None => {
56                    validity.append(false);
57                    values.push(T::default());
58                }
59                Some(e) => {
60                    validity.append(true);
61                    values.push(e);
62                }
63            }
64        }
65        Self::new(values.freeze(), Validity::from(validity.freeze()))
66    }
67
68    pub fn buffer<T: NativePType>(&self) -> Buffer<T> {
69        if T::PTYPE != self.ptype() {
70            vortex_panic!(
71                "Attempted to get buffer of type {} from array of type {}",
72                T::PTYPE,
73                self.ptype()
74            )
75        }
76        Buffer::from_byte_buffer(self.byte_buffer().clone())
77    }
78
79    pub fn into_buffer<T: NativePType>(self) -> Buffer<T> {
80        if T::PTYPE != self.ptype() {
81            vortex_panic!(
82                "Attempted to get buffer of type {} from array of type {}",
83                T::PTYPE,
84                self.ptype()
85            )
86        }
87        Buffer::from_byte_buffer(self.buffer)
88    }
89
90    /// Extract a mutable buffer from the PrimitiveArray. Attempts to do this with zero-copy
91    /// if the buffer is uniquely owned, otherwise will make a copy.
92    pub fn into_buffer_mut<T: NativePType>(self) -> BufferMut<T> {
93        if T::PTYPE != self.ptype() {
94            vortex_panic!(
95                "Attempted to get buffer_mut of type {} from array of type {}",
96                T::PTYPE,
97                self.ptype()
98            )
99        }
100        self.into_buffer()
101            .try_into_mut()
102            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
103    }
104
105    /// Try to extract a mutable buffer from the PrimitiveArray with zero copy.
106    #[allow(clippy::panic_in_result_fn)]
107    pub fn try_into_buffer_mut<T: NativePType>(self) -> Result<BufferMut<T>, PrimitiveArray> {
108        if T::PTYPE != self.ptype() {
109            vortex_panic!(
110                "Attempted to get buffer_mut of type {} from array of type {}",
111                T::PTYPE,
112                self.ptype()
113            )
114        }
115        let validity = self.validity().clone();
116        Buffer::<T>::from_byte_buffer(self.into_byte_buffer())
117            .try_into_mut()
118            .map_err(|buffer| PrimitiveArray::new(buffer, validity))
119    }
120}
121
122impl<T: NativePType> FromIterator<T> for PrimitiveArray {
123    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
124        let values = BufferMut::from_iter(iter);
125        PrimitiveArray::new(values, Validity::NonNullable)
126    }
127}
128
129impl<T: NativePType> IntoArray for Buffer<T> {
130    fn into_array(self) -> ArrayRef {
131        PrimitiveArray::new(self, Validity::NonNullable).into_array()
132    }
133}
134
135impl<T: NativePType> IntoArray for BufferMut<T> {
136    fn into_array(self) -> ArrayRef {
137        self.freeze().into_array()
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use vortex_buffer::BufferMut;
144    use vortex_dtype::{Nullability, PType};
145    use vortex_mask::MaskMut;
146    use vortex_vector::primitive::PVector;
147
148    use super::*;
149
150    #[test]
151    fn test_try_from_vector_with_nulls_nullable() {
152        // Create a vector with some null values: [Some(1), None, Some(3), Some(4), None].
153        let mut values = BufferMut::<i32>::with_capacity(5);
154        values.extend_from_slice(&[1, 0, 3, 4, 0]);
155
156        let mut validity = MaskMut::with_capacity(5);
157        validity.append_n(true, 1);
158        validity.append_n(false, 1);
159        validity.append_n(true, 1);
160        validity.append_n(true, 1);
161        validity.append_n(false, 1);
162
163        let pvector =
164            PVector::try_new(values.freeze(), validity.freeze()).expect("Failed to create PVector");
165
166        // This should succeed since we're allowing nulls.
167        let result =
168            PrimitiveArray::try_from_vector(pvector.into(), Nullability::Nullable).unwrap();
169
170        assert_eq!(result.len(), 5);
171        assert_eq!(result.ptype(), PType::I32);
172        assert!(result.is_valid(0));
173        assert!(!result.is_valid(1));
174        assert!(result.is_valid(2));
175        assert!(result.is_valid(3));
176        assert!(!result.is_valid(4));
177    }
178
179    #[test]
180    fn test_try_from_vector_non_nullable_with_nulls_errors() {
181        // Create a vector with null values: [Some(1), None, Some(3)].
182        let mut values = BufferMut::<i32>::with_capacity(3);
183        values.extend_from_slice(&[1, 0, 3]);
184
185        let mut validity = MaskMut::with_capacity(3);
186        validity.append_n(true, 1);
187        validity.append_n(false, 1);
188        validity.append_n(true, 1);
189
190        let pvector =
191            PVector::try_new(values.freeze(), validity.freeze()).expect("Failed to create PVector");
192
193        // This should fail because we're trying to create a non-nullable array from data with
194        // nulls.
195        let result = PrimitiveArray::try_from_vector(pvector.into(), Nullability::NonNullable);
196
197        assert!(result.is_err());
198        assert!(
199            result
200                .unwrap_err()
201                .to_string()
202                .contains("non-nullable `PrimitiveArray`")
203        );
204    }
205}
vortex_array/arrays/primitive/array/conversion.rs

vortex_array/arrays/primitive/array/
conversion.rs
