vortex_array/arrays/decimal/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::{Itertools, MinMaxResult};
5use vortex_dtype::DecimalDType;
6use vortex_error::VortexExpect;
7use vortex_scalar::{BigCast, DecimalValueType, i256};
8
9use crate::arrays::DecimalArray;
10use crate::vtable::ValidityHelper;
11
12/// Maps a decimal precision into the smallest type that can represent it.
13pub fn smallest_decimal_value_type(decimal_dtype: &DecimalDType) -> DecimalValueType {
14    match decimal_dtype.precision() {
15        1..=2 => DecimalValueType::I8,
16        3..=4 => DecimalValueType::I16,
17        5..=9 => DecimalValueType::I32,
18        10..=18 => DecimalValueType::I64,
19        19..=38 => DecimalValueType::I128,
20        39..=76 => DecimalValueType::I256,
21        0 => unreachable!("precision must be greater than 0"),
22        p => unreachable!("precision larger than 76 is invalid found precision {p}"),
23    }
24}
25
26/// True if `value_type` can represent every value of the type `dtype`.
27pub fn is_compatible_decimal_value_type(value_type: DecimalValueType, dtype: DecimalDType) -> bool {
28    value_type >= smallest_decimal_value_type(&dtype)
29}
30
31macro_rules! try_downcast {
32    ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{
33        // Collect the min/max of the values
34        let minmax = $array.buffer::<$src>().iter().copied().minmax();
35        match minmax {
36            MinMaxResult::NoElements => return $array,
37            MinMaxResult::OneElement(_) => return $array,
38            MinMaxResult::MinMax(min, max) => {
39                $(
40                    if <$dst as BigCast>::from(min).is_some() && <$dst as BigCast>::from(max).is_some() {
41                        return DecimalArray::new::<$dst>(
42                            $array
43                                .buffer::<$src>()
44                                .into_iter()
45                                .map(|v| <$dst as BigCast>::from(v).vortex_expect("decimal conversion failure"))
46                                .collect(),
47                            $array.decimal_dtype(),
48                            $array.validity().clone(),
49                        );
50                    }
51                )*
52
53                return $array;
54            }
55        }
56    }};
57}
58
59/// Attempt to narrow the decimal array to any smaller supported type.
60pub fn narrowed_decimal(decimal_array: DecimalArray) -> DecimalArray {
61    match decimal_array.values_type() {
62        // Cannot narrow any more
63        DecimalValueType::I8 => decimal_array,
64        DecimalValueType::I16 => {
65            try_downcast!(decimal_array, from: i16, to: i8)
66        }
67        DecimalValueType::I32 => {
68            try_downcast!(decimal_array, from: i32, to: i8, i16)
69        }
70        DecimalValueType::I64 => {
71            try_downcast!(decimal_array, from: i64, to: i8, i16, i32)
72        }
73        DecimalValueType::I128 => {
74            try_downcast!(decimal_array, from: i128, to: i8, i16, i32, i64)
75        }
76        DecimalValueType::I256 => {
77            try_downcast!(decimal_array, from: i256, to: i8, i16, i32, i64, i128)
78        }
79        _ => decimal_array,
80    }
81}