Skip to main content

vortex_array/arrays/decimal/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use itertools::MinMaxResult;
6use vortex_buffer::Buffer;
7use vortex_error::VortexExpect;
8
9use crate::arrays::DecimalArray;
10use crate::arrays::decimal::DecimalArrayExt;
11use crate::dtype::DecimalType;
12use crate::dtype::NativeDecimalType;
13use crate::dtype::i256;
14use crate::match_each_decimal_value_type;
15
16/// Return the array's unscaled values widened to `W`, which must be at least as wide as the
17/// array's storage type.
18pub(crate) fn widened_buffer<W: NativeDecimalType>(array: &DecimalArray) -> Buffer<W> {
19    if array.values_type() == W::DECIMAL_TYPE {
20        return array.buffer::<W>();
21    }
22    match_each_decimal_value_type!(array.values_type(), |T| {
23        array
24            .buffer::<T>()
25            .iter()
26            .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed"))
27            .collect()
28    })
29}
30
31macro_rules! try_downcast {
32    ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{
33        use crate::dtype::BigCast;
34
35        // Collect the min/max of the values
36        let minmax = $array.buffer::<$src>().iter().copied().minmax();
37        match minmax {
38            MinMaxResult::NoElements => return $array,
39            MinMaxResult::OneElement(_) => return $array,
40            MinMaxResult::MinMax(min, max) => {
41                $(
42                    if <$dst as BigCast>::from(min).is_some() && <$dst as BigCast>::from(max).is_some() {
43                        return DecimalArray::new::<$dst>(
44                            $array
45                                .buffer::<$src>()
46                                .into_iter()
47                                .map(|v| <$dst as BigCast>::from(v).vortex_expect("decimal conversion failure"))
48                                .collect(),
49                            $array.decimal_dtype(),
50                            $array
51                                .validity()
52                                .vortex_expect("decimal validity should be derivable"),
53                        );
54                    }
55                )*
56
57                return $array;
58            }
59        }
60    }};
61}
62
63/// Attempt to narrow the decimal array to any smaller supported type.
64pub fn narrowed_decimal(decimal_array: DecimalArray) -> DecimalArray {
65    match decimal_array.values_type() {
66        // Cannot narrow any more
67        DecimalType::I8 => decimal_array,
68        DecimalType::I16 => {
69            try_downcast!(decimal_array, from: i16, to: i8)
70        }
71        DecimalType::I32 => {
72            try_downcast!(decimal_array, from: i32, to: i8, i16)
73        }
74        DecimalType::I64 => {
75            try_downcast!(decimal_array, from: i64, to: i8, i16, i32)
76        }
77        DecimalType::I128 => {
78            try_downcast!(decimal_array, from: i128, to: i8, i16, i32, i64)
79        }
80        DecimalType::I256 => {
81            try_downcast!(decimal_array, from: i256, to: i8, i16, i32, i64, i128)
82        }
83    }
84}