vortex_array/builders/
extension.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_ensure;
8use vortex_mask::Mask;
9
10use crate::ArrayRef;
11use crate::ExecutionCtx;
12use crate::IntoArray;
13use crate::arrays::ExtensionArray;
14use crate::arrays::extension::ExtensionArrayExt;
15use crate::builders::ArrayBuilder;
16use crate::builders::DEFAULT_BUILDER_CAPACITY;
17use crate::builders::builder_with_capacity;
18use crate::canonical::Canonical;
19use crate::dtype::DType;
20use crate::dtype::extension::ExtDTypeRef;
21use crate::scalar::ExtScalar;
22use crate::scalar::Scalar;
23
24/// The builder for building a [`ExtensionArray`].
25pub struct ExtensionBuilder {
26    dtype: DType,
27    storage: Box<dyn ArrayBuilder>,
28}
29
30impl ExtensionBuilder {
31    /// Creates a new `ExtensionBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`].
32    pub fn new(ext_dtype: ExtDTypeRef) -> Self {
33        Self::with_capacity(ext_dtype, DEFAULT_BUILDER_CAPACITY)
34    }
35
36    /// Creates a new `ExtensionBuilder` with the given `capacity`.
37    pub fn with_capacity(ext_dtype: ExtDTypeRef, capacity: usize) -> Self {
38        Self {
39            storage: builder_with_capacity(ext_dtype.storage_dtype(), capacity),
40            dtype: DType::Extension(ext_dtype),
41        }
42    }
43
44    /// Appends an extension `value` to the builder.
45    pub fn append_value(&mut self, value: ExtScalar) -> VortexResult<()> {
46        self.storage.append_scalar(&value.to_storage_scalar())
47    }
48
49    /// Appends the values of a canonical [`ExtensionArray`] to the builder by appending its
50    /// storage array to the underlying storage builder.
51    pub(crate) fn append_extension_array(
52        &mut self,
53        array: &ExtensionArray,
54        ctx: &mut ExecutionCtx,
55    ) -> VortexResult<()> {
56        array
57            .storage_array()
58            .append_to_builder(self.storage.as_mut(), ctx)
59    }
60
61    /// Finishes the builder directly into a [`ExtensionArray`].
62    pub fn finish_into_extension(&mut self) -> ExtensionArray {
63        let storage = self.storage.finish();
64        ExtensionArray::new(self.ext_dtype(), storage)
65    }
66
67    /// The [`ExtDType`] of this builder.
68    ///
69    /// [`ExtDType`]: crate::dtype::extension::ExtDType
70    fn ext_dtype(&self) -> ExtDTypeRef {
71        if let DType::Extension(ext_dtype) = &self.dtype {
72            ext_dtype.clone()
73        } else {
74            unreachable!()
75        }
76    }
77}
78
79impl ArrayBuilder for ExtensionBuilder {
80    fn as_any(&self) -> &dyn Any {
81        self
82    }
83
84    fn as_any_mut(&mut self) -> &mut dyn Any {
85        self
86    }
87
88    fn dtype(&self) -> &DType {
89        &self.dtype
90    }
91
92    fn len(&self) -> usize {
93        self.storage.len()
94    }
95
96    fn append_zeros(&mut self, n: usize) {
97        self.storage.append_zeros(n)
98    }
99
100    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
101        self.storage.append_nulls(n)
102    }
103
104    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
105        vortex_ensure!(
106            scalar.dtype() == self.dtype(),
107            "ExtensionBuilder expected scalar with dtype {}, got {}",
108            self.dtype(),
109            scalar.dtype()
110        );
111
112        self.append_value(scalar.as_extension())
113    }
114
115    fn reserve_exact(&mut self, capacity: usize) {
116        self.storage.reserve_exact(capacity)
117    }
118
119    unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
120        unsafe { self.storage.set_validity_unchecked(validity) };
121    }
122
123    fn finish(&mut self) -> ArrayRef {
124        self.finish_into_extension().into_array()
125    }
126
127    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
128        Canonical::Extension(self.finish_into_extension())
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::VortexSessionExecute;
136    use crate::array_session;
137    use crate::arrays::PrimitiveArray;
138    use crate::assert_arrays_eq;
139    use crate::builders::ArrayBuilder;
140    use crate::dtype::Nullability;
141    use crate::extension::datetime::Date;
142    use crate::extension::datetime::TimeUnit;
143    use crate::scalar::Scalar;
144
145    #[test]
146    fn test_append_scalar() {
147        let mut ctx = array_session().create_execution_ctx();
148        let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased();
149
150        let mut builder = ExtensionBuilder::new(ext_dtype.clone());
151
152        // Test appending a valid extension value.
153        let storage1 = Scalar::from(Some(42i32));
154        let ext_scalar1 = Scalar::extension::<Date>(TimeUnit::Days, storage1);
155        builder.append_scalar(&ext_scalar1).unwrap();
156
157        // Test appending another value.
158        let storage2 = Scalar::from(Some(84i32));
159        let ext_scalar2 = Scalar::extension::<Date>(TimeUnit::Days, storage2);
160        builder.append_scalar(&ext_scalar2).unwrap();
161
162        // Test appending null value.
163        let null_storage = Scalar::null(DType::Primitive(
164            crate::dtype::PType::I32,
165            Nullability::Nullable,
166        ));
167        let null_scalar = Scalar::extension::<Date>(TimeUnit::Days, null_storage);
168        builder.append_scalar(&null_scalar).unwrap();
169
170        let array = builder.finish_into_extension();
171        let expected = ExtensionArray::new(
172            ext_dtype.clone(),
173            PrimitiveArray::from_option_iter([Some(42i32), Some(84), None]).into_array(),
174        );
175
176        assert_arrays_eq!(&array, &expected, &mut ctx);
177        assert_eq!(array.len(), 3);
178
179        // Test wrong dtype error.
180        let mut builder = ExtensionBuilder::new(ext_dtype);
181        let wrong_scalar = Scalar::from(true);
182        assert!(builder.append_scalar(&wrong_scalar).is_err());
183    }
184}
vortex_array/builders/extension.rs

vortex_array/builders/
extension.rs
