vortex_array/arrays/decimal/
mod.rs

1mod compute;
2mod macros;
3mod ops;
4mod serde;
5
6use vortex_buffer::{Buffer, ByteBuffer};
7use vortex_dtype::{DType, DecimalDType};
8use vortex_error::{VortexResult, vortex_panic};
9use vortex_scalar::{DecimalValueType, NativeDecimalType};
10
11use crate::builders::ArrayBuilder;
12use crate::stats::{ArrayStats, StatsSetRef};
13use crate::validity::Validity;
14use crate::vtable::{
15    ArrayVTable, CanonicalVTable, NotSupported, VTable, ValidityHelper,
16    ValidityVTableFromValidityHelper, VisitorVTable,
17};
18use crate::{ArrayBufferVisitor, ArrayChildVisitor, Canonical, EncodingId, EncodingRef, vtable};
19
20vtable!(Decimal);
21
22impl VTable for DecimalVTable {
23    type Array = DecimalArray;
24    type Encoding = DecimalEncoding;
25
26    type ArrayVTable = Self;
27    type CanonicalVTable = Self;
28    type OperationsVTable = Self;
29    type ValidityVTable = ValidityVTableFromValidityHelper;
30    type VisitorVTable = Self;
31    type ComputeVTable = NotSupported;
32    type EncodeVTable = NotSupported;
33    type SerdeVTable = Self;
34
35    fn id(_encoding: &Self::Encoding) -> EncodingId {
36        EncodingId::new_ref("vortex.decimal")
37    }
38
39    fn encoding(_array: &Self::Array) -> EncodingRef {
40        EncodingRef::new_ref(DecimalEncoding.as_ref())
41    }
42}
43
44#[derive(Clone, Debug)]
45pub struct DecimalEncoding;
46
47/// Maps a decimal precision into the smallest type that can represent it.
48pub fn smallest_storage_type(decimal_dtype: &DecimalDType) -> DecimalValueType {
49    match decimal_dtype.precision() {
50        1..=2 => DecimalValueType::I8,
51        3..=4 => DecimalValueType::I16,
52        5..=9 => DecimalValueType::I32,
53        10..=18 => DecimalValueType::I64,
54        19..=38 => DecimalValueType::I128,
55        39..=76 => DecimalValueType::I256,
56        0 => unreachable!("precision must be greater than 0"),
57        p => unreachable!("precision larger than 76 is invalid found precision {p}"),
58    }
59}
60
61/// Array for decimal-typed real numbers
62#[derive(Clone, Debug)]
63pub struct DecimalArray {
64    dtype: DType,
65    values: ByteBuffer,
66    values_type: DecimalValueType,
67    validity: Validity,
68    stats_set: ArrayStats,
69}
70
71impl DecimalArray {
72    /// Creates a new [`DecimalArray`] from a [`Buffer`] and [`Validity`], without checking
73    /// any invariants.
74    pub fn new<T: NativeDecimalType>(
75        buffer: Buffer<T>,
76        decimal_dtype: DecimalDType,
77        validity: Validity,
78    ) -> Self {
79        if let Some(len) = validity.maybe_len() {
80            if buffer.len() != len {
81                vortex_panic!(
82                    "Buffer and validity length mismatch: buffer={}, validity={}",
83                    buffer.len(),
84                    len,
85                );
86            }
87        }
88
89        Self {
90            dtype: DType::Decimal(decimal_dtype, validity.nullability()),
91            values: buffer.into_byte_buffer(),
92            values_type: T::VALUES_TYPE,
93            validity,
94            stats_set: ArrayStats::default(),
95        }
96    }
97
98    /// Returns the underlying [`ByteBuffer`] of the array.
99    pub fn byte_buffer(&self) -> ByteBuffer {
100        self.values.clone()
101    }
102
103    pub fn buffer<T: NativeDecimalType>(&self) -> Buffer<T> {
104        if self.values_type != T::VALUES_TYPE {
105            vortex_panic!(
106                "Cannot extract Buffer<{:?}> for DecimalArray with values_type {:?}",
107                T::VALUES_TYPE,
108                self.values_type,
109            );
110        }
111        Buffer::<T>::from_byte_buffer(self.values.clone())
112    }
113
114    /// Returns the decimal type information
115    pub fn decimal_dtype(&self) -> DecimalDType {
116        match &self.dtype {
117            DType::Decimal(decimal_dtype, _) => *decimal_dtype,
118            _ => vortex_panic!("Expected Decimal dtype, got {:?}", self.dtype),
119        }
120    }
121
122    pub fn values_type(&self) -> DecimalValueType {
123        self.values_type
124    }
125
126    pub fn precision(&self) -> u8 {
127        self.decimal_dtype().precision()
128    }
129
130    pub fn scale(&self) -> i8 {
131        self.decimal_dtype().scale()
132    }
133}
134
135impl ArrayVTable<DecimalVTable> for DecimalVTable {
136    fn len(array: &DecimalArray) -> usize {
137        let divisor = match array.values_type {
138            DecimalValueType::I8 => 1,
139            DecimalValueType::I16 => 2,
140            DecimalValueType::I32 => 4,
141            DecimalValueType::I64 => 8,
142            DecimalValueType::I128 => 16,
143            DecimalValueType::I256 => 32,
144            ty => vortex_panic!("unknown decimal value type {:?}", ty),
145        };
146        array.values.len() / divisor
147    }
148
149    fn dtype(array: &DecimalArray) -> &DType {
150        &array.dtype
151    }
152
153    fn stats(array: &DecimalArray) -> StatsSetRef<'_> {
154        array.stats_set.to_ref(array.as_ref())
155    }
156}
157
158impl VisitorVTable<DecimalVTable> for DecimalVTable {
159    fn visit_buffers(array: &DecimalArray, visitor: &mut dyn ArrayBufferVisitor) {
160        visitor.visit_buffer(&array.values);
161    }
162
163    fn visit_children(array: &DecimalArray, visitor: &mut dyn ArrayChildVisitor) {
164        visitor.visit_validity(array.validity(), array.len())
165    }
166}
167
168impl CanonicalVTable<DecimalVTable> for DecimalVTable {
169    fn canonicalize(array: &DecimalArray) -> VortexResult<Canonical> {
170        Ok(Canonical::Decimal(array.clone()))
171    }
172
173    fn append_to_builder(array: &DecimalArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
174        builder.extend_from_array(array.as_ref())
175    }
176}
177
178impl ValidityHelper for DecimalArray {
179    fn validity(&self) -> &Validity {
180        &self.validity
181    }
182}
183
184#[cfg(test)]
185mod test {
186    use arrow_array::Decimal128Array;
187
188    #[test]
189    fn test_decimal() {
190        // They pass it b/c the DType carries the information. No other way to carry a
191        // dtype except via the array.
192        let value = Decimal128Array::new_null(100);
193        let numeric = value.value(10);
194        assert_eq!(numeric, 0i128);
195    }
196}