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

vortex_array/arrays/primitive/array/
conversion.rs
