vortex_array/arrays/decimal/compute/
take.rs1use num_traits::AsPrimitive;
2use vortex_buffer::Buffer;
3use vortex_dtype::{NativePType, match_each_integer_ptype};
4use vortex_error::VortexResult;
5use vortex_scalar::{NativeDecimalType, match_each_decimal_value_type};
6
7use crate::arrays::{DecimalArray, DecimalVTable};
8use crate::compute::{TakeKernel, TakeKernelAdapter};
9use crate::vtable::ValidityHelper;
10use crate::{Array, ArrayRef, ToCanonical, register_kernel};
11
12impl TakeKernel for DecimalVTable {
13 fn take(&self, array: &DecimalArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
14 let indices = indices.to_primitive()?;
15
16 let decimal = match_each_decimal_value_type!(array.values_type(), |$D| {
17 match_each_integer_ptype!(indices.ptype(), |$I| {
18 let buffer = take_to_buffer::<$I, $D>(indices.as_slice::<$I>(), array.buffer::<$D>().as_slice());
19 DecimalArray::new(buffer, array.decimal_dtype(), array.validity().clone())
20 })
21 });
22
23 Ok(decimal.to_array())
24 }
25}
26
27register_kernel!(TakeKernelAdapter(DecimalVTable).lift());
28
29#[inline]
30fn take_to_buffer<I: NativePType + AsPrimitive<usize>, T: NativeDecimalType>(
31 indices: &[I],
32 values: &[T],
33) -> Buffer<T> {
34 indices.iter().map(|idx| values[idx.as_()]).collect()
35}
36
37#[cfg(test)]
38mod tests {
39 use vortex_buffer::buffer;
40 use vortex_dtype::DecimalDType;
41
42 use crate::IntoArray;
43 use crate::arrays::{DecimalArray, DecimalVTable};
44 use crate::compute::take;
45 use crate::validity::Validity;
46
47 #[test]
48 fn test_take() {
49 let array = DecimalArray::new(
50 buffer![10i128, 11i128, 12i128, 13i128],
51 DecimalDType::new(19, 1),
52 Validity::NonNullable,
53 );
54
55 let indices = buffer![0, 2, 3].into_array();
56 let taken = take(array.as_ref(), indices.as_ref()).unwrap();
57 let taken_decimals = taken.as_::<DecimalVTable>();
58 assert_eq!(
59 taken_decimals.buffer::<i128>(),
60 buffer![10i128, 12i128, 13i128]
61 );
62 assert_eq!(taken_decimals.decimal_dtype(), DecimalDType::new(19, 1));
63 }
64}