vortex_array/arrays/decimal/
mod.rs

1mod 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
48/// Maps a decimal precision into the smallest type that can represent it.
49pub 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/// Array for decimal-typed real numbers
63#[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    /// Creates a new [`DecimalArray`] from a [`Buffer`] and [`Validity`], without checking
74    /// any invariants.
75    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    /// Returns the underlying [`ByteBuffer`] of the array.
100    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    /// Returns the decimal type information
116    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        // They pass it b/c the DType carries the information. No other way to carry a
219        // dtype except via the array.
220        let value = Decimal128Array::new_null(100);
221        let numeric = value.value(10);
222        assert_eq!(numeric, 0i128);
223    }
224}
vortex_array/arrays/decimal/mod.rs

vortex_array/arrays/decimal/
mod.rs
