vortex_array/arrays/decimal/compute/
slice.rs

1use vortex_buffer::Buffer;
2use vortex_dtype::DecimalDType;
3use vortex_error::VortexResult;
4
5use crate::arrays::{DecimalArray, DecimalEncoding, NativeDecimalType};
6use crate::compute::SliceFn;
7use crate::validity::Validity;
8use crate::{Array, ArrayRef, match_each_decimal_value_type};
9
10impl SliceFn<&DecimalArray> for DecimalEncoding {
11    fn slice(&self, array: &DecimalArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
12        let sliced = match_each_decimal_value_type!(array.values_type, |$S| {
13            slice_typed(
14                array.buffer::<$S>(),
15                start,
16                stop,
17                array.decimal_dtype(),
18                array.validity.clone(),
19            )
20        });
21
22        Ok(sliced)
23    }
24}
25
26fn slice_typed<T: NativeDecimalType>(
27    values: Buffer<T>,
28    start: usize,
29    end: usize,
30    decimal_dtype: DecimalDType,
31    validity: Validity,
32) -> ArrayRef {
33    let sliced = values.slice(start..end);
34    DecimalArray::new(sliced, decimal_dtype, validity).into_array()
35}
36
37#[cfg(test)]
38mod tests {
39    use vortex_buffer::buffer;
40    use vortex_dtype::DecimalDType;
41
42    use crate::Array;
43    use crate::arrays::DecimalArray;
44    use crate::compute::slice;
45    use crate::validity::Validity;
46
47    #[test]
48    fn test_slice() {
49        let array = DecimalArray::new(
50            buffer![100i128, 200i128, 300i128],
51            DecimalDType::new(3, 2),
52            Validity::NonNullable,
53        );
54
55        let sliced = slice(&array, 1, 3).unwrap();
56        assert_eq!(sliced.len(), 2);
57
58        let decimal = sliced.as_any().downcast_ref::<DecimalArray>().unwrap();
59        assert_eq!(decimal.buffer::<i128>(), buffer![200i128, 300i128]);
60    }
61}