vortex_array/arrays/decimal/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::Buffer;
5use vortex_dtype::DecimalDType;
6use vortex_scalar::{DecimalValue, NativeDecimalType, Scalar, match_each_decimal_value_type};
7
8use crate::arrays::{DecimalArray, DecimalVTable};
9use crate::validity::Validity;
10use crate::vtable::OperationsVTable;
11use crate::{ArrayRef, IntoArray};
12
13impl OperationsVTable<DecimalVTable> for DecimalVTable {
14    fn slice(array: &DecimalArray, start: usize, stop: usize) -> ArrayRef {
15        match_each_decimal_value_type!(array.values_type(), |D| {
16            slice_typed(
17                array.buffer::<D>(),
18                start,
19                stop,
20                array.decimal_dtype(),
21                array.validity.clone(),
22            )
23        })
24    }
25
26    fn scalar_at(array: &DecimalArray, index: usize) -> Scalar {
27        match_each_decimal_value_type!(array.values_type(), |D| {
28            Scalar::decimal(
29                DecimalValue::from(array.buffer::<D>()[index]),
30                array.decimal_dtype(),
31                array.dtype().nullability(),
32            )
33        })
34    }
35}
36
37fn slice_typed<T: NativeDecimalType>(
38    values: Buffer<T>,
39    start: usize,
40    end: usize,
41    decimal_dtype: DecimalDType,
42    validity: Validity,
43) -> ArrayRef {
44    let sliced = values.slice(start..end);
45    let validity = validity.slice(start, end);
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}