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

vortex_array/builders/
extension.rs
