vortex_array/arrays/decimal/compute/
slice.rs1use 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}