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

vortex_array/arrays/decimal/compute/
take.rs
