Skip to main content

vortex_array/arrays/union/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_err;
6use vortex_error::vortex_panic;
7
8use crate::ExecutionCtx;
9use crate::array::ArrayView;
10use crate::array::OperationsVTable;
11use crate::arrays::Union;
12use crate::arrays::union::UnionArrayExt;
13use crate::arrays::union::UnionArraySlotsExt;
14use crate::scalar::Scalar;
15
16impl OperationsVTable<Union> for Union {
17    type ProbeState = ();
18
19    fn scalar_at(
20        array: ArrayView<'_, Union>,
21        index: usize,
22        ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Scalar> {
24        let type_id_scalar = array.type_ids().execute_scalar(index, ctx)?;
25        let Some(type_id) = type_id_scalar.as_primitive().typed_value::<u8>() else {
26            return Ok(Scalar::null(array.dtype().clone()));
27        };
28
29        let Some(child_index) = array.variants().tag_to_child_index(type_id) else {
30            vortex_panic!("Unknown UnionArray type ID {type_id}")
31        };
32        let child = array
33            .child(child_index)
34            .ok_or_else(|| vortex_err!("UnionArray is missing child {child_index}"))?;
35        let child_scalar = child.execute_scalar(index, ctx)?;
36
37        Scalar::union(
38            array.variants().clone(),
39            type_id,
40            child_scalar,
41            array.dtype().nullability(),
42        )
43    }
44}