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