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;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_mask::Mask;
11
12use crate::arrays::DecimalArray;
13use crate::arrays::decimal::DecimalArrayExt;
14use crate::dtype::DecimalType;
15use crate::dtype::NativeDecimalType;
16use crate::dtype::i256;
17use crate::match_each_decimal_value_type;
18
19/// Return the array's unscaled values widened to `W`, which must be at least as wide as the
20/// array's storage type.
21pub(crate) fn widened_buffer<W: NativeDecimalType>(array: &DecimalArray) -> Buffer<W> {
22    if array.values_type() == W::DECIMAL_TYPE {
23        return array.buffer::<W>();
24    }
25    match_each_decimal_value_type!(array.values_type(), |T| {
26        array
27            .buffer::<T>()
28            .iter()
29            .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed"))
30            .collect()
31    })
32}
33
34/// Return the array's unscaled values converted to exactly `W`, whatever the array's
35/// storage type. Zero-copy when the array is already stored at `W`.
36///
37/// Widening is lossless. Narrowing fails for any *valid* value that does not fit `W`.
38/// Null slots may hold arbitrary bytes and never fail; their contents in the returned
39/// buffer are likewise arbitrary and must not be read.
40pub fn converted_buffer<W: NativeDecimalType>(
41    array: &DecimalArray,
42    validity: &Mask,
43) -> VortexResult<Buffer<W>> {
44    // Widening can never fail, so it needs no validation pass.
45    if array.values_type() <= W::DECIMAL_TYPE {
46        return Ok(widened_buffer(array));
47    }
48    match_each_decimal_value_type!(array.values_type(), |T| {
49        let src = array.buffer::<T>();
50        match validity {
51            Mask::AllTrue(_) => {
52                // Keeping the overflow scan branchless and vectorizable. Only on overflow
53                // do we rescan for diagnostics.
54                let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none());
55                if any_overflow {
56                    let (i, v) = src
57                        .iter()
58                        .enumerate()
59                        .find(|&(_, v)| W::from(*v).is_none())
60                        .vortex_expect("overflow scan found an overflowing value");
61                    vortex_bail!(
62                        "decimal value {v} at index {i} does not fit {}",
63                        W::DECIMAL_TYPE
64                    );
65                }
66            }
67            Mask::AllFalse(_) => return Ok(Buffer::zeroed(src.len())),
68            Mask::Values(values) => {
69                let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none());
70                if any_overflow {
71                    for (i, v) in src.iter().enumerate() {
72                        if values.value(i) && W::from(*v).is_none() {
73                            vortex_bail!(
74                                "decimal value {v} at index {i} does not fit {}",
75                                W::DECIMAL_TYPE
76                            );
77                        }
78                    }
79                }
80            }
81        }
82        // The convert pass is infallible: every valid value fits, and out-of-range null-slot
83        // garbage is mapped to the default value. Callers must still ignore null slots.
84        Ok(src
85            .iter()
86            .map(|v| W::from(*v).unwrap_or_default())
87            .collect())
88    })
89}
90
91macro_rules! try_downcast {
92    ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{
93        use crate::dtype::BigCast;
94
95        // Collect the min/max of the values
96        let minmax = $array.buffer::<$src>().iter().copied().minmax();
97        match minmax {
98            MinMaxResult::NoElements => return $array,
99            MinMaxResult::OneElement(_) => return $array,
100            MinMaxResult::MinMax(min, max) => {
101                $(
102                    if <$dst as BigCast>::from(min).is_some() && <$dst as BigCast>::from(max).is_some() {
103                        return DecimalArray::new::<$dst>(
104                            $array
105                                .buffer::<$src>()
106                                .into_iter()
107                                .map(|v| <$dst as BigCast>::from(v).vortex_expect("decimal conversion failure"))
108                                .collect(),
109                            $array.decimal_dtype(),
110                            $array
111                                .validity()
112                                .vortex_expect("decimal validity should be derivable"),
113                        );
114                    }
115                )*
116
117                return $array;
118            }
119        }
120    }};
121}
122
123/// Attempt to narrow the decimal array to any smaller supported type.
124pub fn narrowed_decimal(decimal_array: DecimalArray) -> DecimalArray {
125    match decimal_array.values_type() {
126        // Cannot narrow any more
127        DecimalType::I8 => decimal_array,
128        DecimalType::I16 => {
129            try_downcast!(decimal_array, from: i16, to: i8)
130        }
131        DecimalType::I32 => {
132            try_downcast!(decimal_array, from: i32, to: i8, i16)
133        }
134        DecimalType::I64 => {
135            try_downcast!(decimal_array, from: i64, to: i8, i16, i32)
136        }
137        DecimalType::I128 => {
138            try_downcast!(decimal_array, from: i128, to: i8, i16, i32, i64)
139        }
140        DecimalType::I256 => {
141            try_downcast!(decimal_array, from: i256, to: i8, i16, i32, i64, i128)
142        }
143    }
144}
vortex_array/arrays/decimal/utils.rs

vortex_array/arrays/decimal/
utils.rs
