vortex_runend/compute/
take_from.rs

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

vortex_runend/compute/
take_from.rs
