vortex_array/arrays/decimal/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::Buffer;
5use vortex_error::VortexResult;
6
7use crate::ArrayRef;
8use crate::arrays::DecimalArray;
9use crate::arrays::DecimalVTable;
10use crate::arrays::PrimitiveArray;
11use crate::arrays::TakeExecute;
12use crate::dtype::IntegerPType;
13use crate::dtype::NativeDecimalType;
14use crate::executor::ExecutionCtx;
15use crate::match_each_decimal_value_type;
16use crate::match_each_integer_ptype;
17use crate::vtable::ValidityHelper;
18
19impl TakeExecute for DecimalVTable {
20    fn take(
21        array: &DecimalArray,
22        indices: &ArrayRef,
23        ctx: &mut ExecutionCtx,
24    ) -> VortexResult<Option<ArrayRef>> {
25        let indices = indices.to_array().execute::<PrimitiveArray>(ctx)?;
26        let validity = array.validity().take(&indices.to_array())?;
27
28        // TODO(joe): if the true count of take indices validity is low, only take array values with
29        // valid indices.
30        let decimal = match_each_decimal_value_type!(array.values_type(), |D| {
31            match_each_integer_ptype!(indices.ptype(), |I| {
32                let buffer =
33                    take_to_buffer::<I, D>(indices.as_slice::<I>(), array.buffer::<D>().as_slice());
34                // SAFETY: Take operation preserves decimal dtype and creates valid buffer.
35                // Validity is computed correctly from the parent array and indices.
36                unsafe { DecimalArray::new_unchecked(buffer, array.decimal_dtype(), validity) }
37            })
38        });
39
40        Ok(Some(decimal.to_array()))
41    }
42}
43
44#[inline]
45fn take_to_buffer<I: IntegerPType, T: NativeDecimalType>(indices: &[I], values: &[T]) -> Buffer<T> {
46    indices.iter().map(|idx| values[idx.as_()]).collect()
47}
48
49#[cfg(test)]
50mod tests {
51    use rstest::rstest;
52    use vortex_buffer::Buffer;
53    use vortex_buffer::buffer;
54
55    use crate::IntoArray;
56    use crate::arrays::DecimalArray;
57    use crate::arrays::PrimitiveArray;
58    use crate::assert_arrays_eq;
59    use crate::compute::conformance::take::test_take_conformance;
60    use crate::dtype::DecimalDType;
61    use crate::validity::Validity;
62
63    #[test]
64    fn test_take() {
65        let ddtype = DecimalDType::new(19, 1);
66        let array = DecimalArray::new(
67            buffer![10i128, 11i128, 12i128, 13i128],
68            ddtype,
69            Validity::NonNullable,
70        );
71
72        let indices = buffer![0, 2, 3].into_array();
73        let taken = array.take(indices.to_array()).unwrap();
74
75        let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype);
76        assert_arrays_eq!(expected, taken);
77    }
78
79    #[test]
80    fn test_take_null_indices() {
81        let ddtype = DecimalDType::new(19, 1);
82        let array = DecimalArray::new(
83            buffer![i128::MAX, 11i128, 12i128, 13i128],
84            ddtype,
85            Validity::NonNullable,
86        );
87
88        let indices = PrimitiveArray::from_option_iter([None, Some(2), Some(3)]).into_array();
89        let taken = array.take(indices.to_array()).unwrap();
90
91        let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype);
92        assert_arrays_eq!(expected, taken);
93    }
94
95    #[rstest]
96    #[case(DecimalArray::new(
97        buffer![100i128, 200i128, 300i128, 400i128, 500i128],
98        DecimalDType::new(19, 2),
99        Validity::NonNullable,
100    ))]
101    #[case(DecimalArray::new(
102        buffer![10i64, 20i64, 30i64, 40i64, 50i64],
103        DecimalDType::new(10, 1),
104        Validity::NonNullable,
105    ))]
106    #[case(DecimalArray::new(
107        buffer![1i32, 2i32, 3i32, 4i32, 5i32],
108        DecimalDType::new(5, 0),
109        Validity::NonNullable,
110    ))]
111    #[case(DecimalArray::new(
112        buffer![1000i128, 2000i128, 3000i128, 4000i128, 5000i128],
113        DecimalDType::new(19, 3),
114        Validity::from_iter([true, false, true, true, false]),
115    ))]
116    #[case(DecimalArray::new(
117        buffer![42i128],
118        DecimalDType::new(19, 0),
119        Validity::NonNullable,
120    ))]
121    #[case({
122        let values: Vec<i128> = (0..100).map(|i| i * 1000).collect();
123        DecimalArray::new(
124            Buffer::from_iter(values),
125            DecimalDType::new(19, 4),
126            Validity::NonNullable,
127        )
128    })]
129    fn test_take_decimal_conformance(#[case] array: DecimalArray) {
130        test_take_conformance(&array.to_array());
131    }
132}
vortex_array/arrays/decimal/compute/take.rs

vortex_array/arrays/decimal/compute/
take.rs
