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