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

vortex_array/builders/
extension.rs
