vortex_array/arrays/varbin/compute/
take.rs

1use num_traits::PrimInt;
2use vortex_dtype::{DType, NativePType, match_each_integer_ptype};
3use vortex_error::{VortexResult, vortex_err, vortex_panic};
4use vortex_mask::Mask;
5
6use crate::arrays::VarBinVTable;
7use crate::arrays::varbin::VarBinArray;
8use crate::arrays::varbin::builder::VarBinBuilder;
9use crate::compute::{TakeKernel, TakeKernelAdapter};
10use crate::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
11
12impl TakeKernel for VarBinVTable {
13    fn take(&self, array: &VarBinArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
14        let offsets = array.offsets().to_primitive()?;
15        let data = array.bytes();
16        let indices = indices.to_primitive()?;
17        match_each_integer_ptype!(offsets.ptype(), |O| {
18            match_each_integer_ptype!(indices.ptype(), |I| {
19                Ok(take(
20                    array
21                        .dtype()
22                        .clone()
23                        .union_nullability(indices.dtype().nullability()),
24                    offsets.as_slice::<O>(),
25                    data.as_slice(),
26                    indices.as_slice::<I>(),
27                    array.validity_mask()?,
28                    indices.validity_mask()?,
29                )?
30                .into_array())
31            })
32        })
33    }
34}
35
36register_kernel!(TakeKernelAdapter(VarBinVTable).lift());
37
38fn take<I: NativePType, O: NativePType + PrimInt>(
39    dtype: DType,
40    offsets: &[O],
41    data: &[u8],
42    indices: &[I],
43    validity_mask: Mask,
44    indices_validity_mask: Mask,
45) -> VortexResult<VarBinArray> {
46    if !validity_mask.all_true() || !indices_validity_mask.all_true() {
47        return Ok(take_nullable(
48            dtype,
49            offsets,
50            data,
51            indices,
52            validity_mask,
53            indices_validity_mask,
54        ));
55    }
56
57    let mut builder = VarBinBuilder::<O>::with_capacity(indices.len());
58    for &idx in indices {
59        let idx = idx
60            .to_usize()
61            .ok_or_else(|| vortex_err!("Failed to convert index to usize: {}", idx))?;
62        let start = offsets[idx]
63            .to_usize()
64            .ok_or_else(|| vortex_err!("Failed to convert offset to usize: {}", offsets[idx]))?;
65        let stop = offsets[idx + 1].to_usize().ok_or_else(|| {
66            vortex_err!("Failed to convert offset to usize: {}", offsets[idx + 1])
67        })?;
68        builder.append_value(&data[start..stop]);
69    }
70    Ok(builder.finish(dtype))
71}
72
73fn take_nullable<I: NativePType, O: NativePType + PrimInt>(
74    dtype: DType,
75    offsets: &[O],
76    data: &[u8],
77    indices: &[I],
78    data_validity: Mask,
79    indices_validity: Mask,
80) -> VarBinArray {
81    let mut builder = VarBinBuilder::<O>::with_capacity(indices.len());
82    for (idx, data_idx) in indices.iter().enumerate() {
83        if !indices_validity.value(idx) {
84            builder.append_null();
85            continue;
86        }
87        let data_idx = data_idx
88            .to_usize()
89            .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx));
90        if data_validity.value(data_idx) {
91            let start = offsets[data_idx].to_usize().unwrap_or_else(|| {
92                vortex_panic!("Failed to convert offset to usize: {}", offsets[data_idx])
93            });
94            let stop = offsets[data_idx + 1].to_usize().unwrap_or_else(|| {
95                vortex_panic!(
96                    "Failed to convert offset to usize: {}",
97                    offsets[data_idx + 1]
98                )
99            });
100            builder.append_value(&data[start..stop]);
101        } else {
102            builder.append_null();
103        }
104    }
105    builder.finish(dtype)
106}
107
108#[cfg(test)]
109mod tests {
110    use vortex_dtype::{DType, Nullability};
111
112    use crate::Array;
113    use crate::arrays::{PrimitiveArray, VarBinArray};
114    use crate::compute::take;
115
116    #[test]
117    fn test_null_take() {
118        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
119
120        let idx1: PrimitiveArray = (0..1).collect();
121
122        assert_eq!(
123            take(arr.as_ref(), idx1.as_ref()).unwrap().dtype(),
124            &DType::Utf8(Nullability::NonNullable)
125        );
126
127        let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]);
128
129        assert_eq!(
130            take(arr.as_ref(), idx2.as_ref()).unwrap().dtype(),
131            &DType::Utf8(Nullability::Nullable)
132        );
133    }
134}
vortex_array/arrays/varbin/compute/take.rs

vortex_array/arrays/varbin/compute/
take.rs
