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

vortex_array/builders/
extension.rs
