vortex_array/
vectors.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_dtype::DType;
7use vortex_dtype::NativeDecimalType;
8use vortex_dtype::NativePType;
9use vortex_dtype::Nullability::NonNullable;
10use vortex_dtype::match_each_decimal_value_type;
11use vortex_dtype::match_each_native_ptype;
12use vortex_error::VortexExpect;
13use vortex_vector::Vector;
14use vortex_vector::VectorOps;
15use vortex_vector::binaryview::BinaryViewType;
16use vortex_vector::binaryview::BinaryViewVector;
17use vortex_vector::bool::BoolVector;
18use vortex_vector::decimal::DVector;
19use vortex_vector::decimal::DecimalVector;
20use vortex_vector::fixed_size_list::FixedSizeListVector;
21use vortex_vector::listview::ListViewVector;
22use vortex_vector::null::NullVector;
23use vortex_vector::primitive::PVector;
24use vortex_vector::primitive::PrimitiveVector;
25use vortex_vector::struct_::StructVector;
26
27use crate::ArrayRef;
28use crate::IntoArray;
29use crate::arrays::BoolArray;
30use crate::arrays::DecimalArray;
31use crate::arrays::ExtensionArray;
32use crate::arrays::FixedSizeListArray;
33use crate::arrays::ListViewArray;
34use crate::arrays::NullArray;
35use crate::arrays::PrimitiveArray;
36use crate::arrays::StructArray;
37use crate::arrays::VarBinViewArray;
38use crate::validity::Validity;
39
40/// Trait for converting vector types into arrays.
41pub trait VectorIntoArray {
42    /// Converts the vector into an array of the specified data type.
43    fn into_array(self, dtype: &DType) -> ArrayRef;
44}
45
46impl VectorIntoArray for Vector {
47    fn into_array(self, dtype: &DType) -> ArrayRef {
48        match dtype {
49            DType::Null => self.into_null().into_array(dtype),
50            DType::Bool(_) => self.into_bool().into_array(dtype),
51            DType::Primitive(..) => self.into_primitive().into_array(dtype),
52            DType::Decimal(..) => self.into_decimal().into_array(dtype),
53            DType::Utf8(_) => self.into_string().into_array(dtype),
54            DType::Binary(_) => self.into_binary().into_array(dtype),
55            DType::List(..) => self.into_list().into_array(dtype),
56            DType::FixedSizeList(..) => self.into_fixed_size_list().into_array(dtype),
57            DType::Struct(..) => self.into_struct().into_array(dtype),
58            DType::Extension(ext_dtype) => {
59                let storage = self.into_array(ext_dtype.storage_dtype());
60                ExtensionArray::new(ext_dtype.clone(), storage).into_array()
61            }
62        }
63    }
64}
65
66impl VectorIntoArray for NullVector {
67    fn into_array(self, dtype: &DType) -> ArrayRef {
68        assert!(matches!(dtype, DType::Null));
69        NullArray::new(self.len()).into_array()
70    }
71}
72
73impl VectorIntoArray for BoolVector {
74    fn into_array(self, dtype: &DType) -> ArrayRef {
75        assert!(matches!(dtype, DType::Bool(_)));
76
77        let (bits, validity) = self.into_parts();
78        BoolArray::from_bit_buffer(bits, Validity::from_mask(validity, dtype.nullability()))
79            .into_array()
80    }
81}
82
83impl VectorIntoArray for PrimitiveVector {
84    fn into_array(self, dtype: &DType) -> ArrayRef {
85        match_each_native_ptype!(self.ptype(), |T| {
86            <T as NativePType>::downcast(self).into_array(dtype)
87        })
88    }
89}
90
91impl<T: NativePType> VectorIntoArray for PVector<T> {
92    fn into_array(self, dtype: &DType) -> ArrayRef {
93        assert!(matches!(dtype, DType::Primitive(_, _)));
94        assert_eq!(T::PTYPE, dtype.as_ptype());
95
96        let (values, validity) = self.into_parts();
97        // SAFETY: vectors maintain all invariants required for array creation
98        unsafe {
99            PrimitiveArray::new_unchecked::<T>(
100                values,
101                Validity::from_mask(validity, dtype.nullability()),
102            )
103        }
104        .into_array()
105    }
106}
107
108impl VectorIntoArray for DecimalVector {
109    fn into_array(self, dtype: &DType) -> ArrayRef {
110        match_each_decimal_value_type!(self.decimal_type(), |D| {
111            <D as NativeDecimalType>::downcast(self).into_array(dtype)
112        })
113    }
114}
115
116impl<D: NativeDecimalType> VectorIntoArray for DVector<D> {
117    fn into_array(self, dtype: &DType) -> ArrayRef {
118        assert!(matches!(dtype, DType::Decimal(_, _)));
119
120        let nullability = dtype.nullability();
121        let dec_dtype = dtype
122            .as_decimal_opt()
123            .vortex_expect("expected decimal DType");
124        assert_eq!(dec_dtype.precision(), self.precision());
125        assert_eq!(dec_dtype.scale(), self.scale());
126
127        let (_ps, values, validity) = self.into_parts();
128        // SAFETY: vectors maintain all invariants required for array creation
129        unsafe {
130            DecimalArray::new_unchecked::<D>(
131                values,
132                *dec_dtype,
133                Validity::from_mask(validity, nullability),
134            )
135        }
136        .into_array()
137    }
138}
139
140impl<T: BinaryViewType> VectorIntoArray for BinaryViewVector<T> {
141    fn into_array(self, dtype: &DType) -> ArrayRef {
142        assert!(matches!(dtype, DType::Utf8(_)));
143
144        let (views, buffers, validity) = self.into_parts();
145        let validity = Validity::from_mask(validity, dtype.nullability());
146
147        let buffers = Arc::try_unwrap(buffers).unwrap_or_else(|b| (*b).clone());
148
149        // SAFETY: vectors maintain all invariants required for array creation
150        unsafe {
151            VarBinViewArray::new_unchecked(
152                views,
153                buffers.into_iter().collect(),
154                dtype.clone(),
155                validity,
156            )
157        }
158        .into_array()
159    }
160}
161
162impl VectorIntoArray for ListViewVector {
163    fn into_array(self, dtype: &DType) -> ArrayRef {
164        assert!(matches!(dtype, DType::List(_, _)));
165
166        let (elements, offsets, sizes, validity) = self.into_parts();
167        let validity = Validity::from_mask(validity, dtype.nullability());
168
169        let elements_dtype = dtype.as_list_element_opt().vortex_expect("expected list");
170        let elements = Arc::try_unwrap(elements)
171            .unwrap_or_else(|e| (*e).clone())
172            .into_array(elements_dtype);
173
174        let offsets_dtype = DType::Primitive(offsets.ptype(), NonNullable);
175        let offsets = offsets.into_array(&offsets_dtype);
176
177        let sizes_dtype = DType::Primitive(sizes.ptype(), NonNullable);
178        let sizes = sizes.into_array(&sizes_dtype);
179
180        // SAFETY: vectors maintain all invariants required for array creation
181        unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }.into_array()
182    }
183}
184
185impl VectorIntoArray for FixedSizeListVector {
186    fn into_array(self, dtype: &DType) -> ArrayRef {
187        assert!(matches!(dtype, DType::FixedSizeList(_, _, _)));
188
189        let len = self.len();
190        let (elements, size, validity) = self.into_parts();
191        let validity = Validity::from_mask(validity, dtype.nullability());
192
193        let elements_dtype = dtype
194            .as_fixed_size_list_element_opt()
195            .vortex_expect("expected fixed size list");
196        let elements = Arc::try_unwrap(elements)
197            .unwrap_or_else(|e| (*e).clone())
198            .into_array(elements_dtype);
199
200        // SAFETY: vectors maintain all invariants required for array creation
201        unsafe { FixedSizeListArray::new_unchecked(elements, size, validity, len) }.into_array()
202    }
203}
204
205impl VectorIntoArray for StructVector {
206    fn into_array(self, dtype: &DType) -> ArrayRef {
207        assert!(matches!(dtype, DType::Struct(_, _)));
208
209        let len = self.len();
210        let (fields, validity) = self.into_parts();
211        let validity = Validity::from_mask(validity, dtype.nullability());
212
213        let struct_fields = dtype.as_struct_fields();
214        assert_eq!(fields.len(), struct_fields.nfields());
215
216        let field_arrays: Vec<ArrayRef> = Arc::try_unwrap(fields)
217            .unwrap_or_else(|f| (*f).clone())
218            .into_iter()
219            .zip(struct_fields.fields())
220            .map(|(field_vector, field_dtype)| field_vector.into_array(&field_dtype))
221            .collect();
222
223        // SAFETY: vectors maintain all invariants required for array creation
224        unsafe { StructArray::new_unchecked(field_arrays, struct_fields.clone(), len, validity) }
225            .into_array()
226    }
227}
vortex_array/vectors.rs

vortex_array/
vectors.rs
