vortex_array/builders/
decimal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_buffer::BufferMut;
7use vortex_dtype::{DType, DecimalDType, Nullability};
8use vortex_error::{VortexExpect, vortex_panic};
9use vortex_mask::Mask;
10use vortex_scalar::{BigCast, NativeDecimalType, i256, match_each_decimal_value_type};
11
12use crate::arrays::DecimalArray;
13use crate::builders::{ArrayBuilder, DEFAULT_BUILDER_CAPACITY, LazyNullBufferBuilder};
14use crate::canonical::Canonical;
15use crate::{Array, ArrayRef, IntoArray, ToCanonical};
16
17/// The builder for building a [`DecimalArray`].
18///
19/// The output will be a new [`DecimalArray`] holding values of `T`. Any value that is a valid
20/// [decimal type][NativeDecimalType] can be appended to the builder and it will be immediately
21/// coerced into the target type.
22pub struct DecimalBuilder {
23    dtype: DType,
24    values: DecimalBuffer,
25    nulls: LazyNullBufferBuilder,
26}
27
28/// Wrapper around the typed builder.
29///
30/// We want to be able to downcast a `Box<dyn ArrayBuilder>` to a [`DecimalBuilder`] and we
31/// generally don't have enough type information to get the `T` at the call site, so we instead use
32/// this to hold values and can push values into the correct buffer type generically.
33enum DecimalBuffer {
34    I8(BufferMut<i8>),
35    I16(BufferMut<i16>),
36    I32(BufferMut<i32>),
37    I64(BufferMut<i64>),
38    I128(BufferMut<i128>),
39    I256(BufferMut<i256>),
40}
41
42macro_rules! delegate_fn {
43    ($self:expr, | $tname:ident, $buffer:ident | $body:block) => {{
44        #[allow(unused)]
45        match $self {
46            DecimalBuffer::I8(buffer) => {
47                type $tname = i8;
48                let $buffer = buffer;
49                $body
50            }
51            DecimalBuffer::I16(buffer) => {
52                type $tname = i16;
53                let $buffer = buffer;
54                $body
55            }
56            DecimalBuffer::I32(buffer) => {
57                type $tname = i32;
58                let $buffer = buffer;
59                $body
60            }
61            DecimalBuffer::I64(buffer) => {
62                type $tname = i64;
63                let $buffer = buffer;
64                $body
65            }
66            DecimalBuffer::I128(buffer) => {
67                type $tname = i128;
68                let $buffer = buffer;
69                $body
70            }
71            DecimalBuffer::I256(buffer) => {
72                type $tname = i256;
73                let $buffer = buffer;
74                $body
75            }
76        }
77    }};
78}
79
80impl DecimalBuilder {
81    /// Creates a new `DecimalBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`].
82    pub fn new<T: NativeDecimalType>(precision: u8, scale: i8, nullability: Nullability) -> Self {
83        Self::with_capacity::<T>(
84            DEFAULT_BUILDER_CAPACITY,
85            DecimalDType::new(precision, scale),
86            nullability,
87        )
88    }
89
90    /// Creates a new `DecimalBuilder` with the given `capacity`.
91    pub fn with_capacity<T: NativeDecimalType>(
92        capacity: usize,
93        decimal: DecimalDType,
94        nullability: Nullability,
95    ) -> Self {
96        Self {
97            dtype: DType::Decimal(decimal, nullability),
98            values: match_each_decimal_value_type!(T::VALUES_TYPE, |D| {
99                DecimalBuffer::from(BufferMut::<D>::with_capacity(capacity))
100            }),
101            nulls: LazyNullBufferBuilder::new(capacity),
102        }
103    }
104
105    /// Appends a decimal `value` to the builder.
106    pub fn append_value<V: NativeDecimalType>(&mut self, value: V) {
107        self.values.push(value);
108        self.nulls.append_non_null();
109    }
110
111    /// Appends an optional decimal value to the builder.
112    ///
113    /// If the value is `Some`, it appends the decimal value. If the value is `None`, it appends a
114    /// null.
115    ///
116    /// # Panics
117    ///
118    /// This method will panic if the input is `None` and the builder is non-nullable.
119    pub fn append_option<V: NativeDecimalType>(&mut self, value: Option<V>) {
120        match value {
121            Some(value) => self.append_value(value),
122            None => self.append_null(),
123        }
124    }
125
126    /// Finishes the builder directly into a [`DecimalArray`].
127    pub fn finish_into_decimal(&mut self) -> DecimalArray {
128        let validity = self.nulls.finish_with_nullability(self.dtype.nullability());
129
130        let decimal_dtype = *self.decimal_dtype();
131
132        delegate_fn!(std::mem::take(&mut self.values), |T, values| {
133            DecimalArray::new::<T>(values.freeze(), decimal_dtype, validity)
134        })
135    }
136
137    /// The [`DecimalDType`] of this builder.
138    pub fn decimal_dtype(&self) -> &DecimalDType {
139        let DType::Decimal(decimal_dtype, _) = &self.dtype else {
140            vortex_panic!("`DecimalBuilder` somehow had dtype {}", self.dtype);
141        };
142
143        decimal_dtype
144    }
145}
146
147impl ArrayBuilder for DecimalBuilder {
148    fn as_any(&self) -> &dyn Any {
149        self
150    }
151
152    fn as_any_mut(&mut self) -> &mut dyn Any {
153        self
154    }
155
156    fn dtype(&self) -> &DType {
157        &self.dtype
158    }
159
160    fn len(&self) -> usize {
161        self.values.len()
162    }
163
164    fn append_zeros(&mut self, n: usize) {
165        self.values.push_n(0, n);
166        self.nulls.append_n_non_nulls(n);
167    }
168
169    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
170        self.values.push_n(0, n);
171        self.nulls.append_n_nulls(n);
172    }
173
174    unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
175        let decimal_array = array.to_decimal();
176
177        match_each_decimal_value_type!(decimal_array.values_type(), |D| {
178            // Extends the values buffer from another buffer of type D where D can be coerced to the
179            // builder type.
180            self.values
181                .extend(decimal_array.buffer::<D>().iter().copied());
182        });
183
184        self.nulls
185            .append_validity_mask(decimal_array.validity_mask());
186    }
187
188    fn ensure_capacity(&mut self, capacity: usize) {
189        if capacity > self.values.capacity() {
190            self.values.reserve(capacity - self.values.len());
191            self.nulls.ensure_capacity(capacity);
192        }
193    }
194
195    fn set_validity(&mut self, validity: Mask) {
196        self.nulls = LazyNullBufferBuilder::new(validity.len());
197        self.nulls.append_validity_mask(validity);
198    }
199
200    fn finish(&mut self) -> ArrayRef {
201        self.finish_into_decimal().into_array()
202    }
203
204    fn finish_into_canonical(&mut self) -> Canonical {
205        Canonical::Decimal(self.finish_into_decimal())
206    }
207}
208
209impl DecimalBuffer {
210    fn push<V: NativeDecimalType>(&mut self, value: V) {
211        delegate_fn!(self, |T, buffer| {
212            buffer.push(<T as BigCast>::from(value).vortex_expect("decimal conversion failure"))
213        });
214    }
215
216    fn push_n<V: NativeDecimalType>(&mut self, value: V, n: usize) {
217        delegate_fn!(self, |T, buffer| {
218            buffer.push_n(
219                <T as BigCast>::from(value).vortex_expect("decimal conversion failure"),
220                n,
221            )
222        });
223    }
224
225    fn reserve(&mut self, additional: usize) {
226        delegate_fn!(self, |T, buffer| { buffer.reserve(additional) })
227    }
228
229    fn capacity(&self) -> usize {
230        delegate_fn!(self, |T, buffer| { buffer.capacity() })
231    }
232
233    fn len(&self) -> usize {
234        delegate_fn!(self, |T, buffer| { buffer.len() })
235    }
236
237    pub fn extend<I, V: NativeDecimalType>(&mut self, iter: I)
238    where
239        I: Iterator<Item = V>,
240    {
241        delegate_fn!(self, |T, buffer| {
242            buffer.extend(
243                iter.map(|x| <T as BigCast>::from(x).vortex_expect("decimal conversion failure")),
244            )
245        })
246    }
247}
248
249macro_rules! impl_from_buffer {
250    ($T:ty, $variant:ident) => {
251        impl From<BufferMut<$T>> for DecimalBuffer {
252            fn from(buffer: BufferMut<$T>) -> Self {
253                Self::$variant(buffer)
254            }
255        }
256    };
257}
258
259impl_from_buffer!(i8, I8);
260impl_from_buffer!(i16, I16);
261impl_from_buffer!(i32, I32);
262impl_from_buffer!(i64, I64);
263impl_from_buffer!(i128, I128);
264impl_from_buffer!(i256, I256);
265
266impl Default for DecimalBuffer {
267    fn default() -> Self {
268        Self::I8(BufferMut::<i8>::empty())
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use crate::builders::{ArrayBuilder, DecimalBuilder};
275
276    #[test]
277    fn test_mixed_extend() {
278        let values = 42i8;
279
280        let mut i8s = DecimalBuilder::new::<i8>(2, 1, false.into());
281        for v in 0..values {
282            i8s.append_value(v);
283        }
284        let i8s = i8s.finish();
285
286        let mut i128s = DecimalBuilder::new::<i128>(2, 1, false.into());
287        i128s.extend_from_array(&i8s);
288        let i128s = i128s.finish();
289
290        for i in 0..i8s.len() {
291            assert_eq!(i8s.scalar_at(i), i128s.scalar_at(i));
292        }
293    }
294}
vortex_array/builders/decimal.rs

vortex_array/builders/
decimal.rs
