vortex_runend/compute/
take_from.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::Array;
5use vortex_array::ArrayRef;
6use vortex_array::IntoArray;
7use vortex_array::compute::TakeFromKernel;
8use vortex_array::compute::TakeFromKernelAdapter;
9use vortex_array::compute::take;
10use vortex_array::register_kernel;
11use vortex_dtype::DType;
12use vortex_error::VortexResult;
13
14use crate::RunEndArray;
15use crate::RunEndVTable;
16
17impl TakeFromKernel for RunEndVTable {
18    /// Takes values from the source array using run-end encoded indices.
19    ///
20    /// # Arguments
21    ///
22    /// * `indices` - Run-end encoded indices
23    /// * `source` - Array to take values from
24    ///
25    /// # Returns
26    ///
27    /// * `Ok(Some(source))` - If successful
28    /// * `Ok(None)` - If the source array has an unsupported dtype
29    ///
30    fn take_from(
31        &self,
32        indices: &RunEndArray,
33        source: &dyn Array,
34    ) -> VortexResult<Option<ArrayRef>> {
35        // Only `Primitive` and `Bool` are valid run-end value types. - TODO: Support additional DTypes
36        if !matches!(source.dtype(), DType::Primitive(_, _) | DType::Bool(_)) {
37            return Ok(None);
38        }
39
40        // Transform the run-end encoding from storing indices to storing values
41        // by taking values from `source` at positions specified by `indices.values()`.
42        let values = take(source, indices.values())?;
43
44        // Create a new run-end array containing values as values, instead of indices as values.
45        // SAFETY: we are copying ends from an existing valid RunEndArray
46        let ree_array = unsafe {
47            RunEndArray::new_unchecked(
48                indices.ends().clone(),
49                values,
50                indices.offset(),
51                indices.len(),
52            )
53        };
54
55        Ok(Some(ree_array.into_array()))
56    }
57}
58
59register_kernel!(TakeFromKernelAdapter(RunEndVTable).lift());
vortex_runend/compute/take_from.rs

vortex_runend/compute/
take_from.rs
