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

vortex_array/builders/
mod.rs
