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;
9use vortex_mask::Mask;
10use vortex_scalar::ExtScalar;
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    /// Appends an optional extension value to the builder.
43    ///
44    /// If the value is `Some`, it appends the extension value. If the value is `None`, it appends a
45    /// null.
46    ///
47    /// # Panics
48    ///
49    /// This method will panic if the input is `None` and the builder is non-nullable.
50    pub fn append_option(&mut self, value: Option<ExtScalar>) -> VortexResult<()> {
51        match value {
52            Some(value) => self.append_value(value),
53            None => {
54                self.append_nulls(1);
55                Ok(())
56            }
57        }
58    }
59
60    /// Finishes the builder directly into a [`ExtensionArray`].
61    pub fn finish_into_extension(&mut self) -> ExtensionArray {
62        let storage = self.storage.finish();
63        ExtensionArray::new(self.ext_dtype(), storage)
64    }
65
66    /// The [`ExtDType`] of this builder.
67    fn ext_dtype(&self) -> Arc<ExtDType> {
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    unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
102        let ext_array = array.to_extension();
103        self.storage.extend_from_array(ext_array.storage())
104    }
105
106    fn ensure_capacity(&mut self, capacity: usize) {
107        self.storage.ensure_capacity(capacity)
108    }
109
110    fn set_validity(&mut self, validity: Mask) {
111        self.storage.set_validity(validity);
112    }
113
114    fn finish(&mut self) -> ArrayRef {
115        self.finish_into_extension().into_array()
116    }
117
118    fn finish_into_canonical(&mut self) -> Canonical {
119        Canonical::Extension(self.finish_into_extension())
120    }
121}
vortex_array/builders/extension.rs

vortex_array/builders/
extension.rs
