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<T> {
42    /// Converts the vector into an array of the specified data type.
43    fn into_array(self, dtype: &DType) -> T;
44}
45
46impl VectorIntoArray<ArrayRef> for Vector {
47    fn into_array(self, dtype: &DType) -> ArrayRef {
48        match dtype {
49            DType::Null => self.into_null().into_array(dtype).into_array(),
50            DType::Bool(_) => self.into_bool().into_array(dtype).into_array(),
51            DType::Primitive(..) => self.into_primitive().into_array(dtype).into_array(),
52            DType::Decimal(..) => self.into_decimal().into_array(dtype).into_array(),
53            DType::Utf8(_) => self.into_string().into_array(dtype).into_array(),
54            DType::Binary(_) => self.into_binary().into_array(dtype).into_array(),
55            DType::List(..) => self.into_list().into_array(dtype).into_array(),
56            DType::FixedSizeList(..) => self.into_fixed_size_list().into_array(dtype).into_array(),
57            DType::Struct(..) => self.into_struct().into_array(dtype).into_array(),
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<NullArray> for NullVector {
67    fn into_array(self, dtype: &DType) -> NullArray {
68        assert!(matches!(dtype, DType::Null));
69        NullArray::new(self.len())
70    }
71}
72
73impl VectorIntoArray<BoolArray> for BoolVector {
74    fn into_array(self, dtype: &DType) -> BoolArray {
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    }
80}
81
82impl VectorIntoArray<PrimitiveArray> for PrimitiveVector {
83    fn into_array(self, dtype: &DType) -> PrimitiveArray {
84        match_each_native_ptype!(self.ptype(), |T| {
85            <T as NativePType>::downcast(self).into_array(dtype)
86        })
87    }
88}
89
90impl<T: NativePType> VectorIntoArray<PrimitiveArray> for PVector<T> {
91    fn into_array(self, dtype: &DType) -> PrimitiveArray {
92        assert!(matches!(dtype, DType::Primitive(_, _)));
93        assert_eq!(T::PTYPE, dtype.as_ptype());
94
95        let (values, validity) = self.into_parts();
96        // SAFETY: vectors maintain all invariants required for array creation
97        unsafe {
98            PrimitiveArray::new_unchecked::<T>(
99                values,
100                Validity::from_mask(validity, dtype.nullability()),
101            )
102        }
103    }
104}
105
106impl VectorIntoArray<DecimalArray> for DecimalVector {
107    fn into_array(self, dtype: &DType) -> DecimalArray {
108        match_each_decimal_value_type!(self.decimal_type(), |D| {
109            <D as NativeDecimalType>::downcast(self).into_array(dtype)
110        })
111    }
112}
113
114impl<D: NativeDecimalType> VectorIntoArray<DecimalArray> for DVector<D> {
115    fn into_array(self, dtype: &DType) -> DecimalArray {
116        assert!(matches!(dtype, DType::Decimal(_, _)));
117
118        let nullability = dtype.nullability();
119        let dec_dtype = dtype
120            .as_decimal_opt()
121            .vortex_expect("expected decimal DType");
122        assert_eq!(dec_dtype.precision(), self.precision());
123        assert_eq!(dec_dtype.scale(), self.scale());
124
125        let (_ps, values, validity) = self.into_parts();
126        // SAFETY: vectors maintain all invariants required for array creation
127        unsafe {
128            DecimalArray::new_unchecked::<D>(
129                values,
130                *dec_dtype,
131                Validity::from_mask(validity, nullability),
132            )
133        }
134    }
135}
136
137impl<T: BinaryViewType> VectorIntoArray<VarBinViewArray> for BinaryViewVector<T> {
138    fn into_array(self, dtype: &DType) -> VarBinViewArray {
139        assert!(matches!(dtype, DType::Utf8(_)));
140
141        let (views, buffers, validity) = self.into_parts();
142        let validity = Validity::from_mask(validity, dtype.nullability());
143
144        let buffers = Arc::try_unwrap(buffers).unwrap_or_else(|b| (*b).clone());
145
146        // SAFETY: vectors maintain all invariants required for array creation
147        unsafe {
148            VarBinViewArray::new_unchecked(
149                views,
150                buffers.into_iter().collect(),
151                dtype.clone(),
152                validity,
153            )
154        }
155    }
156}
157
158impl VectorIntoArray<ListViewArray> for ListViewVector {
159    fn into_array(self, dtype: &DType) -> ListViewArray {
160        assert!(matches!(dtype, DType::List(_, _)));
161
162        let (elements, offsets, sizes, validity) = self.into_parts();
163        let validity = Validity::from_mask(validity, dtype.nullability());
164
165        let elements_dtype = dtype.as_list_element_opt().vortex_expect("expected list");
166        let elements = Arc::try_unwrap(elements)
167            .unwrap_or_else(|e| (*e).clone())
168            .into_array(elements_dtype);
169
170        let offsets_dtype = DType::Primitive(offsets.ptype(), NonNullable);
171        let offsets = offsets.into_array(&offsets_dtype);
172
173        let sizes_dtype = DType::Primitive(sizes.ptype(), NonNullable);
174        let sizes = sizes.into_array(&sizes_dtype);
175
176        // SAFETY: vectors maintain all invariants required for array creation
177        unsafe {
178            ListViewArray::new_unchecked(
179                elements,
180                offsets.into_array(),
181                sizes.into_array(),
182                validity,
183            )
184        }
185    }
186}
187
188impl VectorIntoArray<FixedSizeListArray> for FixedSizeListVector {
189    fn into_array(self, dtype: &DType) -> FixedSizeListArray {
190        assert!(matches!(dtype, DType::FixedSizeList(_, _, _)));
191
192        let len = self.len();
193        let (elements, size, validity) = self.into_parts();
194        let validity = Validity::from_mask(validity, dtype.nullability());
195
196        let elements_dtype = dtype
197            .as_fixed_size_list_element_opt()
198            .vortex_expect("expected fixed size list");
199        let elements = Arc::try_unwrap(elements)
200            .unwrap_or_else(|e| (*e).clone())
201            .into_array(elements_dtype);
202
203        // SAFETY: vectors maintain all invariants required for array creation
204        unsafe { FixedSizeListArray::new_unchecked(elements, size, validity, len) }
205    }
206}
207
208impl VectorIntoArray<StructArray> for StructVector {
209    fn into_array(self, dtype: &DType) -> StructArray {
210        assert!(matches!(dtype, DType::Struct(_, _)));
211
212        let len = self.len();
213        let (fields, validity) = self.into_parts();
214        let validity = Validity::from_mask(validity, dtype.nullability());
215
216        let struct_fields = dtype.as_struct_fields();
217        assert_eq!(fields.len(), struct_fields.nfields());
218
219        let field_arrays: Vec<ArrayRef> = Arc::try_unwrap(fields)
220            .unwrap_or_else(|f| (*f).clone())
221            .into_iter()
222            .zip(struct_fields.fields())
223            .map(|(field_vector, field_dtype)| field_vector.into_array(&field_dtype))
224            .collect();
225
226        // SAFETY: vectors maintain all invariants required for array creation
227        unsafe { StructArray::new_unchecked(field_arrays, struct_fields.clone(), len, validity) }
228    }
229}
vortex_array/vectors.rs

vortex_array/
vectors.rs
