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    fn ext_dtype(&self) -> ExtDTypeRef {
58        if let DType::Extension(ext_dtype) = &self.dtype {
59            ext_dtype.clone()
60        } else {
61            unreachable!()
62        }
63    }
64}
65
66impl ArrayBuilder for ExtensionBuilder {
67    fn as_any(&self) -> &dyn Any {
68        self
69    }
70
71    fn as_any_mut(&mut self) -> &mut dyn Any {
72        self
73    }
74
75    fn dtype(&self) -> &DType {
76        &self.dtype
77    }
78
79    fn len(&self) -> usize {
80        self.storage.len()
81    }
82
83    fn append_zeros(&mut self, n: usize) {
84        self.storage.append_zeros(n)
85    }
86
87    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
88        self.storage.append_nulls(n)
89    }
90
91    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
92        vortex_ensure!(
93            scalar.dtype() == self.dtype(),
94            "ExtensionBuilder expected scalar with dtype {}, got {}",
95            self.dtype(),
96            scalar.dtype()
97        );
98
99        self.append_value(scalar.as_extension())
100    }
101
102    unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) {
103        #[expect(deprecated)]
104        let ext_array = array.to_extension();
105        self.storage.extend_from_array(ext_array.storage_array())
106    }
107
108    fn reserve_exact(&mut self, capacity: usize) {
109        self.storage.reserve_exact(capacity)
110    }
111
112    unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
113        unsafe { self.storage.set_validity_unchecked(validity) };
114    }
115
116    fn finish(&mut self) -> ArrayRef {
117        self.finish_into_extension().into_array()
118    }
119
120    fn finish_into_canonical(&mut self) -> Canonical {
121        Canonical::Extension(self.finish_into_extension())
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use crate::arrays::PrimitiveArray;
129    use crate::assert_arrays_eq;
130    use crate::builders::ArrayBuilder;
131    use crate::dtype::Nullability;
132    use crate::extension::datetime::Date;
133    use crate::extension::datetime::TimeUnit;
134    use crate::scalar::Scalar;
135
136    #[test]
137    fn test_append_scalar() {
138        let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased();
139
140        let mut builder = ExtensionBuilder::new(ext_dtype.clone());
141
142        // Test appending a valid extension value.
143        let storage1 = Scalar::from(Some(42i32));
144        let ext_scalar1 = Scalar::extension::<Date>(TimeUnit::Days, storage1);
145        builder.append_scalar(&ext_scalar1).unwrap();
146
147        // Test appending another value.
148        let storage2 = Scalar::from(Some(84i32));
149        let ext_scalar2 = Scalar::extension::<Date>(TimeUnit::Days, storage2);
150        builder.append_scalar(&ext_scalar2).unwrap();
151
152        // Test appending null value.
153        let null_storage = Scalar::null(DType::Primitive(
154            crate::dtype::PType::I32,
155            Nullability::Nullable,
156        ));
157        let null_scalar = Scalar::extension::<Date>(TimeUnit::Days, null_storage);
158        builder.append_scalar(&null_scalar).unwrap();
159
160        let array = builder.finish_into_extension();
161        let expected = ExtensionArray::new(
162            ext_dtype.clone(),
163            PrimitiveArray::from_option_iter([Some(42i32), Some(84), None]).into_array(),
164        );
165
166        assert_arrays_eq!(&array, &expected);
167        assert_eq!(array.len(), 3);
168
169        // Test wrong dtype error.
170        let mut builder = ExtensionBuilder::new(ext_dtype);
171        let wrong_scalar = Scalar::from(true);
172        assert!(builder.append_scalar(&wrong_scalar).is_err());
173    }
174}
vortex_array/builders/extension.rs

vortex_array/builders/
extension.rs
