vortex_array/arrays/decimal/
ops.rs

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