vortex_array/arrays/decimal/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_buffer::Buffer;
7use vortex_dtype::DecimalDType;
8use vortex_scalar::{DecimalValue, NativeDecimalType, Scalar, match_each_decimal_value_type};
9
10use crate::arrays::{DecimalArray, DecimalVTable};
11use crate::validity::Validity;
12use crate::vtable::OperationsVTable;
13use crate::{ArrayRef, IntoArray};
14
15impl OperationsVTable<DecimalVTable> for DecimalVTable {
16    fn slice(array: &DecimalArray, range: Range<usize>) -> ArrayRef {
17        match_each_decimal_value_type!(array.values_type(), |D| {
18            slice_typed(
19                array.buffer::<D>(),
20                range,
21                array.decimal_dtype(),
22                array.validity.clone(),
23            )
24        })
25    }
26
27    fn scalar_at(array: &DecimalArray, index: usize) -> Scalar {
28        match_each_decimal_value_type!(array.values_type(), |D| {
29            Scalar::decimal(
30                DecimalValue::from(array.buffer::<D>()[index]),
31                array.decimal_dtype(),
32                array.dtype().nullability(),
33            )
34        })
35    }
36}
37
38fn slice_typed<T: NativeDecimalType>(
39    values: Buffer<T>,
40    range: Range<usize>,
41    decimal_dtype: DecimalDType,
42    validity: Validity,
43) -> ArrayRef {
44    let sliced = values.slice(range.clone());
45    let validity = validity.slice(range);
46    DecimalArray::new(sliced, decimal_dtype, validity).into_array()
47}
48
49#[cfg(test)]
50mod tests {
51    use vortex_buffer::buffer;
52    use vortex_dtype::{DecimalDType, Nullability};
53    use vortex_scalar::{DecimalValue, Scalar};
54
55    use crate::Array;
56    use crate::arrays::{DecimalArray, DecimalVTable};
57    use crate::validity::Validity;
58
59    #[test]
60    fn test_slice() {
61        let array = DecimalArray::new(
62            buffer![100i128, 200i128, 300i128, 4000i128],
63            DecimalDType::new(3, 2),
64            Validity::NonNullable,
65        )
66        .to_array();
67
68        let sliced = array.slice(1..3);
69        assert_eq!(sliced.len(), 2);
70
71        let decimal = sliced.as_::<DecimalVTable>();
72        assert_eq!(decimal.buffer::<i128>(), buffer![200i128, 300i128]);
73    }
74
75    #[test]
76    fn test_slice_nullable() {
77        let array = DecimalArray::new(
78            buffer![100i128, 200i128, 300i128, 4000i128],
79            DecimalDType::new(3, 2),
80            Validity::from_iter([false, true, false, true]),
81        )
82        .to_array();
83
84        let sliced = array.slice(1..3);
85        assert_eq!(sliced.len(), 2);
86    }
87
88    #[test]
89    fn test_scalar_at() {
90        let array = DecimalArray::new(
91            buffer![100i128],
92            DecimalDType::new(3, 2),
93            Validity::NonNullable,
94        );
95
96        assert_eq!(
97            array.scalar_at(0),
98            Scalar::decimal(
99                DecimalValue::I128(100),
100                DecimalDType::new(3, 2),
101                Nullability::NonNullable
102            )
103        );
104    }
105}