vortex_array/builders/
mod.rs

1mod bool;
2mod extension;
3mod lazy_validity_builder;
4mod list;
5mod null;
6mod primitive;
7mod struct_;
8mod varbinview;
9
10use std::any::Any;
11
12pub use bool::*;
13pub use extension::*;
14pub use list::*;
15pub use null::*;
16pub use primitive::*;
17pub use varbinview::*;
18use vortex_dtype::{DType, match_each_native_ptype};
19use vortex_error::{VortexResult, vortex_bail, vortex_err};
20use vortex_scalar::{
21    BinaryScalar, BoolScalar, ExtScalar, ListScalar, PrimitiveScalar, Scalar, ScalarValue,
22    StructScalar, Utf8Scalar,
23};
24
25use crate::builders::struct_::StructBuilder;
26use crate::{Array, ArrayRef};
27
28pub trait ArrayBuilder: Send {
29    fn as_any(&self) -> &dyn Any;
30
31    fn as_any_mut(&mut self) -> &mut dyn Any;
32
33    fn dtype(&self) -> &DType;
34
35    fn len(&self) -> usize;
36
37    fn is_empty(&self) -> bool {
38        self.len() == 0
39    }
40
41    /// Append a "zero" value to the array.
42    fn append_zero(&mut self) {
43        self.append_zeros(1)
44    }
45
46    /// Appends n "zero" values to the array.
47    fn append_zeros(&mut self, n: usize);
48
49    /// Append a "null" value to the array.
50    fn append_null(&mut self) {
51        self.append_nulls(1)
52    }
53
54    /// Appends n "null" values to the array.
55    fn append_nulls(&mut self, n: usize);
56
57    /// Extends the array with the provided array, canonicalizing if necessary.
58    fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()>;
59
60    /// Constructs an Array from the builder components.
61    ///
62    /// # Panics
63    ///
64    /// This function may panic if the builder's methods are called with invalid arguments. If only
65    /// the methods on this interface are used, the builder should not panic. However, specific
66    /// builders have interfaces that may be misued. For example, if the number of values in a
67    /// [PrimitiveBuilder]'s [vortex_buffer::BufferMut] does not match the number of validity bits,
68    /// the PrimitiveBuilder's [Self::finish] will panic.
69    fn finish(&mut self) -> ArrayRef;
70}
71
72pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBuilder> {
73    match dtype {
74        DType::Null => Box::new(NullBuilder::new()),
75        DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity)),
76        DType::Primitive(ptype, n) => {
77            match_each_native_ptype!(ptype, |$P| {
78                Box::new(PrimitiveBuilder::<$P>::with_capacity(*n, capacity))
79            })
80        }
81        DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity(DType::Utf8(*n), capacity)),
82        DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity(
83            DType::Binary(*n),
84            capacity,
85        )),
86        DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity(
87            struct_dtype.clone(),
88            *n,
89            capacity,
90        )),
91        DType::List(dtype, n) => Box::new(ListBuilder::<u64>::with_capacity(
92            dtype.clone(),
93            *n,
94            capacity,
95        )),
96        DType::Extension(ext_dtype) => {
97            Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity))
98        }
99    }
100}
101
102pub trait ArrayBuilderExt: ArrayBuilder {
103    /// A generic function to append a scalar value to the builder.
104    fn append_scalar_value(&mut self, value: ScalarValue) -> VortexResult<()> {
105        if value.is_null() {
106            self.append_null();
107            Ok(())
108        } else {
109            self.append_scalar(&Scalar::new(self.dtype().clone(), value))
110        }
111    }
112
113    /// A generic function to append a scalar to the builder.
114    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
115        if scalar.dtype() != self.dtype() {
116            vortex_bail!(
117                "Builder has dtype {:?}, scalar has {:?}",
118                self.dtype(),
119                scalar.dtype()
120            )
121        }
122        match scalar.dtype() {
123            DType::Null => self
124                .as_any_mut()
125                .downcast_mut::<NullBuilder>()
126                .ok_or_else(|| vortex_err!("Cannot append null scalar to non-null builder"))?
127                .append_null(),
128            DType::Bool(_) => self
129                .as_any_mut()
130                .downcast_mut::<BoolBuilder>()
131                .ok_or_else(|| vortex_err!("Cannot append bool scalar to non-bool builder"))?
132                .append_option(BoolScalar::try_from(scalar)?.value()),
133            DType::Primitive(ptype, ..) => {
134                match_each_native_ptype!(ptype, |$P| {
135                    self
136                    .as_any_mut()
137                    .downcast_mut::<PrimitiveBuilder<$P>>()
138                    .ok_or_else(|| {
139                        vortex_err!("Cannot append primitive scalar to non-primitive builder")
140                    })?
141                    .append_option(PrimitiveScalar::try_from(scalar)?.typed_value::<$P>())
142                })
143            }
144            DType::Utf8(_) => self
145                .as_any_mut()
146                .downcast_mut::<VarBinViewBuilder>()
147                .ok_or_else(|| vortex_err!("Cannot append utf8 scalar to non-utf8 builder"))?
148                .append_option(Utf8Scalar::try_from(scalar)?.value()),
149            DType::Binary(_) => self
150                .as_any_mut()
151                .downcast_mut::<VarBinViewBuilder>()
152                .ok_or_else(|| vortex_err!("Cannot append binary scalar to non-binary builder"))?
153                .append_option(BinaryScalar::try_from(scalar)?.value()),
154            DType::Struct(..) => self
155                .as_any_mut()
156                .downcast_mut::<StructBuilder>()
157                .ok_or_else(|| vortex_err!("Cannot append struct scalar to non-struct builder"))?
158                .append_value(StructScalar::try_from(scalar)?)?,
159            DType::List(..) => self
160                .as_any_mut()
161                .downcast_mut::<ListBuilder<u64>>()
162                .ok_or_else(|| vortex_err!("Cannot append list scalar to non-list builder"))?
163                .append_value(ListScalar::try_from(scalar)?)?,
164            DType::Extension(..) => self
165                .as_any_mut()
166                .downcast_mut::<ExtensionBuilder>()
167                .ok_or_else(|| {
168                    vortex_err!("Cannot append extension scalar to non-extension builder")
169                })?
170                .append_value(ExtScalar::try_from(scalar)?)?,
171        }
172        Ok(())
173    }
174}
175
176impl<T: ?Sized + ArrayBuilder> ArrayBuilderExt for T {}