vortex_array/arrays/union/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::Union;
11use crate::arrays::UnionArray;
12use crate::arrays::dict::TakeReduce;
13use crate::arrays::union::UnionArrayExt;
14use crate::arrays::union::UnionArraySlotsExt;
15use crate::builtins::ArrayBuiltins;
16use crate::scalar::Scalar;
17
18/// Gathers the type IDs and every sparse child at `indices`.
19///
20/// Sparse children are row-aligned with the union, so a gather must visit all of them. Take costs
21/// `O(variants * indices)`, which only the dense encoding fixes.
22///
23/// The type IDs carry the union's validity, so gathering them with the original `indices` turns a
24/// null index into an outer union null. The children are gathered with the nulls filled in, which
25/// keeps their declared variant dtypes.
26impl TakeReduce for Union {
27    fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
28        // An empty union has no row to point at, so the indices must be all null.
29        if array.is_empty() {
30            return UnionArray::constant(&Scalar::null(array.dtype().as_nullable()), indices.len())
31                .map(UnionArray::into_array)
32                .map(Some);
33        }
34
35        let type_ids = array.type_ids().take(indices.clone())?;
36
37        // This stays a lazy node, so the fill runs once per child. `TakeReduce` has no
38        // `ExecutionCtx` to materialize it with, and the cost scales with the indices, not the
39        // data behind them.
40        let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable());
41        let child_indices = indices.clone().fill_null(fill_scalar)?;
42
43        let children: Vec<ArrayRef> = array
44            .iter_children()
45            .map(|child| child.take(child_indices.clone()))
46            .try_collect()?;
47
48        UnionArray::try_new(type_ids, array.variants().clone(), children)
49            .map(UnionArray::into_array)
50            .map(Some)
51    }
52}
vortex_array/arrays/union/compute/take.rs

vortex_array/arrays/union/compute/
take.rs
