vortex_array/builders/
extension.rs

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

vortex_array/builders/
extension.rs
