vortex_array/arrays/decimal/vtable/
operations.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_dtype::NativeDecimalType;
9use vortex_dtype::match_each_decimal_value_type;
10use vortex_scalar::DecimalValue;
11use vortex_scalar::Scalar;
12
13use crate::ArrayRef;
14use crate::IntoArray;
15use crate::arrays::DecimalArray;
16use crate::arrays::DecimalVTable;
17use crate::validity::Validity;
18use crate::vtable::OperationsVTable;
19
20impl OperationsVTable<DecimalVTable> for DecimalVTable {
21    fn slice(array: &DecimalArray, range: Range<usize>) -> ArrayRef {
22        match_each_decimal_value_type!(array.values_type(), |D| {
23            slice_typed(
24                array.buffer::<D>(),
25                range,
26                array.decimal_dtype(),
27                array.validity.clone(),
28            )
29        })
30    }
31
32    fn scalar_at(array: &DecimalArray, index: usize) -> Scalar {
33        match_each_decimal_value_type!(array.values_type(), |D| {
34            Scalar::decimal(
35                DecimalValue::from(array.buffer::<D>()[index]),
36                array.decimal_dtype(),
37                array.dtype().nullability(),
38            )
39        })
40    }
41}
42
43fn slice_typed<T: NativeDecimalType>(
44    values: Buffer<T>,
45    range: Range<usize>,
46    decimal_dtype: DecimalDType,
47    validity: Validity,
48) -> ArrayRef {
49    let sliced = values.slice(range.clone());
50    let validity = validity.slice(range);
51    // SAFETY: Slicing preserves all DecimalArray invariants:
52    // - Buffer is correctly typed and sized from the slice operation.
53    // - Decimal dtype is preserved from the parent array.
54    // - Validity is correctly sliced to match the new length.
55    unsafe { DecimalArray::new_unchecked(sliced, decimal_dtype, validity) }.into_array()
56}
57
58#[cfg(test)]
59mod tests {
60    use vortex_buffer::buffer;
61    use vortex_dtype::DecimalDType;
62    use vortex_dtype::Nullability;
63    use vortex_scalar::DecimalValue;
64    use vortex_scalar::Scalar;
65
66    use crate::Array;
67    use crate::arrays::DecimalArray;
68    use crate::arrays::DecimalVTable;
69    use crate::validity::Validity;
70
71    #[test]
72    fn test_slice() {
73        let array = DecimalArray::new(
74            buffer![100i128, 200i128, 300i128, 4000i128],
75            DecimalDType::new(3, 2),
76            Validity::NonNullable,
77        )
78        .to_array();
79
80        let sliced = array.slice(1..3);
81        assert_eq!(sliced.len(), 2);
82
83        let decimal = sliced.as_::<DecimalVTable>();
84        assert_eq!(decimal.buffer::<i128>(), buffer![200i128, 300i128]);
85    }
86
87    #[test]
88    fn test_slice_nullable() {
89        let array = DecimalArray::new(
90            buffer![100i128, 200i128, 300i128, 4000i128],
91            DecimalDType::new(3, 2),
92            Validity::from_iter([false, true, false, true]),
93        )
94        .to_array();
95
96        let sliced = array.slice(1..3);
97        assert_eq!(sliced.len(), 2);
98    }
99
100    #[test]
101    fn test_scalar_at() {
102        let array = DecimalArray::new(
103            buffer![100i128],
104            DecimalDType::new(3, 2),
105            Validity::NonNullable,
106        );
107
108        assert_eq!(
109            array.scalar_at(0),
110            Scalar::decimal(
111                DecimalValue::I128(100),
112                DecimalDType::new(3, 2),
113                Nullability::NonNullable
114            )
115        );
116    }
117}