vortex_array/arrays/decimal/
mod.rs

1mod compute;
2mod serde;
3
4use vortex_buffer::{Buffer, ByteBuffer};
5use vortex_dtype::{DType, DecimalDType};
6use vortex_error::{VortexResult, vortex_panic};
7use vortex_mask::Mask;
8use vortex_scalar::i256;
9
10use crate::array::{Array, ArrayCanonicalImpl, ArrayValidityImpl, ArrayVariantsImpl};
11use crate::arrays::decimal::serde::DecimalMetadata;
12use crate::builders::ArrayBuilder;
13use crate::stats::{ArrayStats, StatsSetRef};
14use crate::validity::Validity;
15use crate::variants::DecimalArrayTrait;
16use crate::vtable::VTableRef;
17use crate::{
18    ArrayBufferVisitor, ArrayChildVisitor, ArrayImpl, ArrayRef, ArrayStatisticsImpl,
19    ArrayVisitorImpl, Canonical, Encoding, ProstMetadata, try_from_array_ref,
20};
21
22#[derive(Debug)]
23pub struct DecimalEncoding;
24
25pub use crate::arrays::decimal::serde::DecimalValueType;
26
27impl Encoding for DecimalEncoding {
28    type Array = DecimalArray;
29    type Metadata = ProstMetadata<DecimalMetadata>;
30}
31
32/// Type of decimal scalar values.
33pub trait NativeDecimalType: Copy + Eq + Ord {
34    const VALUES_TYPE: DecimalValueType;
35}
36
37impl NativeDecimalType for i8 {
38    const VALUES_TYPE: DecimalValueType = DecimalValueType::I8;
39}
40
41impl NativeDecimalType for i16 {
42    const VALUES_TYPE: DecimalValueType = DecimalValueType::I16;
43}
44
45impl NativeDecimalType for i32 {
46    const VALUES_TYPE: DecimalValueType = DecimalValueType::I32;
47}
48
49impl NativeDecimalType for i64 {
50    const VALUES_TYPE: DecimalValueType = DecimalValueType::I64;
51}
52
53impl NativeDecimalType for i128 {
54    const VALUES_TYPE: DecimalValueType = DecimalValueType::I128;
55}
56
57impl NativeDecimalType for i256 {
58    const VALUES_TYPE: DecimalValueType = DecimalValueType::I256;
59}
60
61/// Maps a decimal precision into the small type that can represent it.
62pub fn precision_to_storage_size(decimal_dtype: &DecimalDType) -> DecimalValueType {
63    match decimal_dtype.precision() {
64        1..=2 => DecimalValueType::I8,
65        3..=4 => DecimalValueType::I16,
66        5..=9 => DecimalValueType::I32,
67        10..=18 => DecimalValueType::I64,
68        19..=38 => DecimalValueType::I128,
69        39..=76 => DecimalValueType::I256,
70        0 => unreachable!("precision must be greater than 0"),
71        p => unreachable!("precision larger than 76 is invalid found precision {p}"),
72    }
73}
74
75/// Array for decimal-typed real numbers
76#[derive(Clone, Debug)]
77pub struct DecimalArray {
78    dtype: DType,
79    values: ByteBuffer,
80    values_type: DecimalValueType,
81    validity: Validity,
82    stats_set: ArrayStats,
83}
84
85try_from_array_ref!(DecimalArray);
86
87impl DecimalArray {
88    /// Creates a new [`DecimalArray`] from a [`Buffer`] and [`Validity`], without checking
89    /// any invariants.
90    pub fn new<T: NativeDecimalType>(
91        buffer: Buffer<T>,
92        decimal_dtype: DecimalDType,
93        validity: Validity,
94    ) -> Self {
95        if let Some(len) = validity.maybe_len() {
96            if buffer.len() != len {
97                vortex_panic!(
98                    "Buffer and validity length mismatch: buffer={}, validity={}",
99                    buffer.len(),
100                    len
101                );
102            }
103        }
104
105        Self {
106            dtype: DType::Decimal(decimal_dtype, validity.nullability()),
107            values: buffer.into_byte_buffer(),
108            values_type: T::VALUES_TYPE,
109            validity,
110            stats_set: ArrayStats::default(),
111        }
112    }
113
114    /// Returns the underlying [`ByteBuffer`] of the array.
115    pub fn byte_buffer(&self) -> ByteBuffer {
116        self.values.clone()
117    }
118
119    pub fn buffer<T: NativeDecimalType>(&self) -> Buffer<T> {
120        if self.values_type != T::VALUES_TYPE {
121            vortex_panic!(
122                "Cannot extract Buffer<{:?}> for DecimalArray with values_type {:?}",
123                T::VALUES_TYPE,
124                self.values_type
125            );
126        }
127        Buffer::<T>::from_byte_buffer(self.values.clone())
128    }
129
130    /// Returns the underlying [`Validity`] of the array.
131    pub fn validity(&self) -> &Validity {
132        &self.validity
133    }
134
135    /// Returns the decimal type information
136    pub fn decimal_dtype(&self) -> DecimalDType {
137        match &self.dtype {
138            DType::Decimal(decimal_dtype, _) => *decimal_dtype,
139            _ => vortex_panic!("Expected Decimal dtype, got {:?}", self.dtype),
140        }
141    }
142
143    pub fn values_type(&self) -> DecimalValueType {
144        self.values_type
145    }
146
147    pub fn precision(&self) -> u8 {
148        self.decimal_dtype().precision()
149    }
150
151    pub fn scale(&self) -> i8 {
152        self.decimal_dtype().scale()
153    }
154}
155
156impl ArrayVisitorImpl<ProstMetadata<DecimalMetadata>> for DecimalArray {
157    fn _metadata(&self) -> ProstMetadata<DecimalMetadata> {
158        ProstMetadata(DecimalMetadata {
159            values_type: self.values_type as i32,
160        })
161    }
162
163    fn _visit_buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
164        visitor.visit_buffer(&self.values);
165    }
166
167    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
168        visitor.visit_validity(self.validity(), self.len())
169    }
170}
171
172impl ArrayImpl for DecimalArray {
173    type Encoding = DecimalEncoding;
174
175    #[inline]
176    fn _len(&self) -> usize {
177        let divisor = match self.values_type {
178            DecimalValueType::I8 => 1,
179            DecimalValueType::I16 => 2,
180            DecimalValueType::I32 => 4,
181            DecimalValueType::I64 => 8,
182            DecimalValueType::I128 => 16,
183            DecimalValueType::I256 => 32,
184        };
185        self.values.len() / divisor
186    }
187
188    #[inline]
189    fn _dtype(&self) -> &DType {
190        &self.dtype
191    }
192
193    #[inline]
194    fn _vtable(&self) -> VTableRef {
195        VTableRef::new_ref(&DecimalEncoding)
196    }
197
198    fn _with_children(&self, _children: &[ArrayRef]) -> VortexResult<Self> {
199        // No non-validity child arrays to replace.
200        Ok(self.clone())
201    }
202}
203
204impl ArrayStatisticsImpl for DecimalArray {
205    fn _stats_ref(&self) -> StatsSetRef<'_> {
206        self.stats_set.to_ref(self)
207    }
208}
209
210impl ArrayVariantsImpl for DecimalArray {
211    fn _as_decimal_typed(&self) -> Option<&dyn DecimalArrayTrait> {
212        Some(self)
213    }
214}
215
216impl DecimalArrayTrait for DecimalArray {}
217
218impl ArrayCanonicalImpl for DecimalArray {
219    #[inline]
220    fn _to_canonical(&self) -> VortexResult<Canonical> {
221        Ok(Canonical::Decimal(self.clone()))
222    }
223
224    #[inline]
225    fn _append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
226        builder.extend_from_array(self)
227    }
228}
229
230impl ArrayValidityImpl for DecimalArray {
231    #[inline]
232    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
233        self.validity.is_valid(index)
234    }
235
236    #[inline]
237    fn _all_valid(&self) -> VortexResult<bool> {
238        self.validity.all_valid()
239    }
240
241    #[inline]
242    fn _all_invalid(&self) -> VortexResult<bool> {
243        self.validity.all_invalid()
244    }
245
246    #[inline]
247    fn _validity_mask(&self) -> VortexResult<Mask> {
248        self.validity.to_mask(self.len())
249    }
250}
251
252#[cfg(test)]
253mod test {
254    use arrow_array::Decimal128Array;
255
256    #[test]
257    fn test_decimal() {
258        // They pass it b/c the DType carries the information. No other way to carry a
259        // dtype except via the array.
260        let value = Decimal128Array::new_null(100);
261        let numeric = value.value(10);
262        assert_eq!(numeric, 0i128);
263    }
264}
vortex_array/arrays/decimal/mod.rs

vortex_array/arrays/decimal/
mod.rs
