vortex_array/builders/
mod.rs1mod 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 fn append_zero(&mut self) {
43 self.append_zeros(1)
44 }
45
46 fn append_zeros(&mut self, n: usize);
48
49 fn append_null(&mut self) {
51 self.append_nulls(1)
52 }
53
54 fn append_nulls(&mut self, n: usize);
56
57 fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()>;
59
60 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 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 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 {}