vortex_runend/compute/
take_from.rs

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