vortex_array/builders/
extension.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use vortex_dtype::{DType, ExtDType};
5use vortex_error::{VortexResult, vortex_bail};
6use vortex_mask::Mask;
7use vortex_scalar::ExtScalar;
8
9use crate::arrays::ExtensionArray;
10use crate::builders::{ArrayBuilder, ArrayBuilderExt, builder_with_capacity};
11use crate::{Array, ArrayRef, Canonical};
12
13pub struct ExtensionBuilder {
14    storage: Box<dyn ArrayBuilder>,
15    dtype: DType,
16}
17
18impl ExtensionBuilder {
19    pub fn new(ext_dtype: Arc<ExtDType>) -> Self {
20        Self::with_capacity(ext_dtype, 1024)
21    }
22
23    pub fn with_capacity(ext_dtype: Arc<ExtDType>, capacity: usize) -> Self {
24        Self {
25            storage: builder_with_capacity(ext_dtype.storage_dtype(), capacity),
26            dtype: DType::Extension(ext_dtype),
27        }
28    }
29
30    pub fn append_value(&mut self, value: ExtScalar) -> VortexResult<()> {
31        self.storage.append_scalar(&value.storage())
32    }
33
34    pub fn append_option(&mut self, value: Option<ExtScalar>) -> VortexResult<()> {
35        match value {
36            Some(value) => self.append_value(value),
37            None => {
38                self.append_nulls(1);
39                Ok(())
40            }
41        }
42    }
43
44    fn ext_dtype(&self) -> Arc<ExtDType> {
45        if let DType::Extension(ext_dtype) = &self.dtype {
46            ext_dtype.clone()
47        } else {
48            unreachable!()
49        }
50    }
51}
52
53impl ArrayBuilder for ExtensionBuilder {
54    fn as_any(&self) -> &dyn Any {
55        self
56    }
57
58    fn as_any_mut(&mut self) -> &mut dyn Any {
59        self
60    }
61
62    fn dtype(&self) -> &DType {
63        &self.dtype
64    }
65
66    fn len(&self) -> usize {
67        self.storage.len()
68    }
69
70    fn append_zeros(&mut self, n: usize) {
71        self.storage.append_zeros(n)
72    }
73
74    fn append_nulls(&mut self, n: usize) {
75        self.storage.append_nulls(n)
76    }
77
78    fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
79        let array = array.to_canonical()?;
80        let Canonical::Extension(array) = array else {
81            vortex_bail!("Expected Extension array, got {:?}", array);
82        };
83        array.storage().append_to_builder(self.storage.as_mut())
84    }
85
86    fn ensure_capacity(&mut self, capacity: usize) {
87        self.storage.ensure_capacity(capacity)
88    }
89
90    fn set_validity(&mut self, validity: Mask) {
91        self.storage.set_validity(validity);
92    }
93
94    fn finish(&mut self) -> ArrayRef {
95        let storage = self.storage.finish();
96        ExtensionArray::new(self.ext_dtype(), storage).into_array()
97    }
98}
vortex_array/builders/extension.rs

vortex_array/builders/
extension.rs
