vortex_array/builders/
decimal.rs

1use std::any::Any;
2
3use vortex_buffer::BufferMut;
4use vortex_dtype::{DType, DecimalDType, Nullability};
5use vortex_error::{VortexResult, vortex_bail, vortex_panic};
6use vortex_mask::Mask;
7
8use crate::arrays::{BoolArray, DecimalArray, NativeDecimalType};
9use crate::builders::ArrayBuilder;
10use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
11use crate::validity::Validity;
12use crate::{Array, ArrayRef, ToCanonical};
13
14pub struct DecimalBuilder<T> {
15    values: BufferMut<T>,
16    nulls: LazyNullBufferBuilder,
17    dtype: DType,
18}
19
20const DEFAULT_BUILDER_CAPACITY: usize = 1024;
21
22impl<T: NativeDecimalType> DecimalBuilder<T> {
23    pub fn new(precision: u8, scale: i8, nullability: Nullability) -> Self {
24        Self::with_capacity(DEFAULT_BUILDER_CAPACITY, precision, scale, nullability)
25    }
26
27    pub fn with_capacity(
28        capacity: usize,
29        precision: u8,
30        scale: i8,
31        nullability: Nullability,
32    ) -> Self {
33        Self {
34            values: BufferMut::with_capacity(capacity),
35            nulls: LazyNullBufferBuilder::new(capacity),
36            dtype: DType::Decimal(DecimalDType::new(precision, scale), nullability),
37        }
38    }
39
40    /// Append a `Mask` to the null buffer.
41    pub fn append_mask(&mut self, mask: Mask) {
42        self.nulls.append_validity_mask(mask);
43    }
44
45    fn extend_with_validity_mask(&mut self, validity_mask: Mask) {
46        self.nulls.append_validity_mask(validity_mask);
47    }
48
49    pub fn append_value(&mut self, value: T) {
50        self.values.push(value);
51        self.nulls.append(true);
52    }
53
54    pub fn append_option(&mut self, value: Option<T>)
55    where
56        Self: Send,
57        T: Default + Send + 'static,
58    {
59        match value {
60            Some(value) => {
61                self.values.push(value);
62                self.nulls.append(true);
63            }
64            None => self.append_null(),
65        }
66    }
67
68    pub fn values(&self) -> &[T] {
69        self.values.as_ref()
70    }
71
72    pub fn finish_into_decimal(&mut self) -> DecimalArray {
73        let nulls = self.nulls.finish();
74
75        if let Some(null_buf) = nulls.as_ref() {
76            assert_eq!(
77                null_buf.len(),
78                self.values.len(),
79                "null buffer length must equal value buffer length"
80            );
81        }
82
83        let validity = match (nulls, self.dtype.nullability()) {
84            (None, Nullability::NonNullable) => Validity::NonNullable,
85            (Some(_), Nullability::NonNullable) => {
86                vortex_panic!("Non-nullable builder has null values")
87            }
88            (None, Nullability::Nullable) => Validity::AllValid,
89            (Some(nulls), Nullability::Nullable) => {
90                if nulls.null_count() == nulls.len() {
91                    Validity::AllInvalid
92                } else {
93                    Validity::Array(BoolArray::from(nulls.into_inner()).into_array())
94                }
95            }
96        };
97
98        let DType::Decimal(decimal_dtype, _) = self.dtype else {
99            vortex_panic!("DecimalBuilder must have Decimal DType");
100        };
101
102        DecimalArray::new(
103            std::mem::take(&mut self.values).freeze(),
104            decimal_dtype,
105            validity,
106        )
107    }
108}
109
110impl<T: NativeDecimalType + Default + Send + 'static> ArrayBuilder for DecimalBuilder<T> {
111    fn as_any(&self) -> &dyn Any {
112        self
113    }
114
115    fn as_any_mut(&mut self) -> &mut dyn Any {
116        self
117    }
118
119    fn dtype(&self) -> &DType {
120        &self.dtype
121    }
122
123    fn len(&self) -> usize {
124        self.values.len()
125    }
126
127    fn append_zeros(&mut self, n: usize) {
128        self.values.push_n(T::default(), n);
129        self.nulls.append_n_non_nulls(n);
130    }
131
132    fn append_nulls(&mut self, n: usize) {
133        self.values.push_n(T::default(), n);
134        self.nulls.append_n_nulls(n);
135    }
136
137    fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
138        let array = array.to_decimal()?;
139
140        let DType::Decimal(decimal_dtype, _) = self.dtype else {
141            vortex_panic!("DecimalBuilder must have Decimal DType");
142        };
143
144        if array.decimal_dtype() != decimal_dtype {
145            vortex_bail!(
146                "Cannot extend from array with different decimal type: {:?} != {:?}",
147                array.decimal_dtype(),
148                decimal_dtype
149            );
150        }
151
152        self.values
153            .extend_from_slice(array.buffer::<T>().as_slice());
154        self.extend_with_validity_mask(array.validity_mask()?);
155
156        Ok(())
157    }
158
159    fn ensure_capacity(&mut self, capacity: usize) {
160        if capacity > self.values.capacity() {
161            self.values.reserve(capacity - self.values.len());
162            self.nulls.ensure_capacity(capacity);
163        }
164    }
165
166    fn set_validity(&mut self, validity: Mask) {
167        self.nulls = LazyNullBufferBuilder::new(validity.len());
168        self.nulls.append_validity_mask(validity);
169    }
170
171    fn finish(&mut self) -> ArrayRef {
172        self.finish_into_decimal().into_array()
173    }
174}
vortex_array/builders/decimal.rs

vortex_array/builders/
decimal.rs
