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    fn scalar_at(
18        array: ArrayView<'_, Union>,
19        index: usize,
20        ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Scalar> {
22        let type_id_scalar = array.type_ids().execute_scalar(index, ctx)?;
23        let Some(type_id) = type_id_scalar.as_primitive().typed_value::<u8>() else {
24            return Ok(Scalar::null(array.dtype().clone()));
25        };
26
27        let Some(child_index) = array.variants().tag_to_child_index(type_id) else {
28            vortex_panic!("Unknown UnionArray type ID {type_id}")
29        };
30        let child = array
31            .child(child_index)
32            .ok_or_else(|| vortex_err!("UnionArray is missing child {child_index}"))?;
33        let child_scalar = child.execute_scalar(index, ctx)?;
34
35        Scalar::union(
36            array.variants().clone(),
37            type_id,
38            child_scalar,
39            array.dtype().nullability(),
40        )
41    }
42}