vortex_array/arrays/decimal/
mod.rs1mod compute;
2mod macros;
3mod ops;
4mod serde;
5
6use arrow_buffer::BooleanBufferBuilder;
7use vortex_buffer::{Buffer, BufferMut, ByteBuffer};
8use vortex_dtype::{DType, DecimalDType};
9use vortex_error::{VortexResult, vortex_panic};
10use vortex_scalar::{DecimalValueType, NativeDecimalType};
11
12use crate::builders::ArrayBuilder;
13use crate::stats::{ArrayStats, StatsSetRef};
14use crate::validity::Validity;
15use crate::vtable::{
16 ArrayVTable, CanonicalVTable, NotSupported, VTable, ValidityHelper,
17 ValidityVTableFromValidityHelper, VisitorVTable,
18};
19use crate::{ArrayBufferVisitor, ArrayChildVisitor, Canonical, EncodingId, EncodingRef, vtable};
20
21vtable!(Decimal);
22
23impl VTable for DecimalVTable {
24 type Array = DecimalArray;
25 type Encoding = DecimalEncoding;
26
27 type ArrayVTable = Self;
28 type CanonicalVTable = Self;
29 type OperationsVTable = Self;
30 type ValidityVTable = ValidityVTableFromValidityHelper;
31 type VisitorVTable = Self;
32 type ComputeVTable = NotSupported;
33 type EncodeVTable = NotSupported;
34 type SerdeVTable = Self;
35
36 fn id(_encoding: &Self::Encoding) -> EncodingId {
37 EncodingId::new_ref("vortex.decimal")
38 }
39
40 fn encoding(_array: &Self::Array) -> EncodingRef {
41 EncodingRef::new_ref(DecimalEncoding.as_ref())
42 }
43}
44
45#[derive(Clone, Debug)]
46pub struct DecimalEncoding;
47
48pub fn smallest_storage_type(decimal_dtype: &DecimalDType) -> DecimalValueType {
50 match decimal_dtype.precision() {
51 1..=2 => DecimalValueType::I8,
52 3..=4 => DecimalValueType::I16,
53 5..=9 => DecimalValueType::I32,
54 10..=18 => DecimalValueType::I64,
55 19..=38 => DecimalValueType::I128,
56 39..=76 => DecimalValueType::I256,
57 0 => unreachable!("precision must be greater than 0"),
58 p => unreachable!("precision larger than 76 is invalid found precision {p}"),
59 }
60}
61
62#[derive(Clone, Debug)]
64pub struct DecimalArray {
65 dtype: DType,
66 values: ByteBuffer,
67 values_type: DecimalValueType,
68 validity: Validity,
69 stats_set: ArrayStats,
70}
71
72impl DecimalArray {
73 pub fn new<T: NativeDecimalType>(
76 buffer: Buffer<T>,
77 decimal_dtype: DecimalDType,
78 validity: Validity,
79 ) -> Self {
80 if let Some(len) = validity.maybe_len() {
81 if buffer.len() != len {
82 vortex_panic!(
83 "Buffer and validity length mismatch: buffer={}, validity={}",
84 buffer.len(),
85 len,
86 );
87 }
88 }
89
90 Self {
91 dtype: DType::Decimal(decimal_dtype, validity.nullability()),
92 values: buffer.into_byte_buffer(),
93 values_type: T::VALUES_TYPE,
94 validity,
95 stats_set: ArrayStats::default(),
96 }
97 }
98
99 pub fn byte_buffer(&self) -> ByteBuffer {
101 self.values.clone()
102 }
103
104 pub fn buffer<T: NativeDecimalType>(&self) -> Buffer<T> {
105 if self.values_type != T::VALUES_TYPE {
106 vortex_panic!(
107 "Cannot extract Buffer<{:?}> for DecimalArray with values_type {:?}",
108 T::VALUES_TYPE,
109 self.values_type,
110 );
111 }
112 Buffer::<T>::from_byte_buffer(self.values.clone())
113 }
114
115 pub fn decimal_dtype(&self) -> DecimalDType {
117 match &self.dtype {
118 DType::Decimal(decimal_dtype, _) => *decimal_dtype,
119 _ => vortex_panic!("Expected Decimal dtype, got {:?}", self.dtype),
120 }
121 }
122
123 pub fn values_type(&self) -> DecimalValueType {
124 self.values_type
125 }
126
127 pub fn precision(&self) -> u8 {
128 self.decimal_dtype().precision()
129 }
130
131 pub fn scale(&self) -> i8 {
132 self.decimal_dtype().scale()
133 }
134
135 pub fn from_option_iter<T: NativeDecimalType, I: IntoIterator<Item = Option<T>>>(
136 iter: I,
137 decimal_dtype: DecimalDType,
138 ) -> Self {
139 let iter = iter.into_iter();
140 let mut values = BufferMut::with_capacity(iter.size_hint().0);
141 let mut validity = BooleanBufferBuilder::new(values.capacity());
142
143 for i in iter {
144 match i {
145 None => {
146 validity.append(false);
147 values.push(T::default());
148 }
149 Some(e) => {
150 validity.append(true);
151 values.push(e);
152 }
153 }
154 }
155 Self::new(
156 values.freeze(),
157 decimal_dtype,
158 Validity::from(validity.finish()),
159 )
160 }
161}
162
163impl ArrayVTable<DecimalVTable> for DecimalVTable {
164 fn len(array: &DecimalArray) -> usize {
165 let divisor = match array.values_type {
166 DecimalValueType::I8 => 1,
167 DecimalValueType::I16 => 2,
168 DecimalValueType::I32 => 4,
169 DecimalValueType::I64 => 8,
170 DecimalValueType::I128 => 16,
171 DecimalValueType::I256 => 32,
172 ty => vortex_panic!("unknown decimal value type {:?}", ty),
173 };
174 array.values.len() / divisor
175 }
176
177 fn dtype(array: &DecimalArray) -> &DType {
178 &array.dtype
179 }
180
181 fn stats(array: &DecimalArray) -> StatsSetRef<'_> {
182 array.stats_set.to_ref(array.as_ref())
183 }
184}
185
186impl VisitorVTable<DecimalVTable> for DecimalVTable {
187 fn visit_buffers(array: &DecimalArray, visitor: &mut dyn ArrayBufferVisitor) {
188 visitor.visit_buffer(&array.values);
189 }
190
191 fn visit_children(array: &DecimalArray, visitor: &mut dyn ArrayChildVisitor) {
192 visitor.visit_validity(array.validity(), array.len())
193 }
194}
195
196impl CanonicalVTable<DecimalVTable> for DecimalVTable {
197 fn canonicalize(array: &DecimalArray) -> VortexResult<Canonical> {
198 Ok(Canonical::Decimal(array.clone()))
199 }
200
201 fn append_to_builder(array: &DecimalArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
202 builder.extend_from_array(array.as_ref())
203 }
204}
205
206impl ValidityHelper for DecimalArray {
207 fn validity(&self) -> &Validity {
208 &self.validity
209 }
210}
211
212#[cfg(test)]
213mod test {
214 use arrow_array::Decimal128Array;
215
216 #[test]
217 fn test_decimal() {
218 let value = Decimal128Array::new_null(100);
221 let numeric = value.value(10);
222 assert_eq!(numeric, 0i128);
223 }
224}