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}