vortex_array/arrays/union/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_error::vortex_ensure;
7use vortex_error::vortex_ensure_eq;
8use vortex_error::vortex_err;
9use vortex_error::vortex_panic;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::ExecutionResult;
16use crate::array::Array;
17use crate::array::ArrayId;
18use crate::array::ArrayParts;
19use crate::array::ArrayView;
20use crate::array::EmptyArrayData;
21use crate::array::VTable;
22use crate::array::with_empty_buffers;
23use crate::arrays::union::UnionArrayExt;
24use crate::arrays::union::array::CHILDREN_OFFSET;
25use crate::arrays::union::array::TYPE_IDS_SLOT;
26use crate::arrays::union::array::make_union_parts;
27use crate::arrays::union::compute::rules::PARENT_RULES;
28use crate::arrays::union::union_type_ids_dtype;
29use crate::buffer::BufferHandle;
30use crate::builders::ArrayBuilder;
31use crate::dtype::DType;
32use crate::serde::ArrayChildren;
33
34mod operations;
35mod validate;
36mod validity;
37
38/// A canonical [`Union`]-encoded array.
39///
40/// This is similar to Arrow's "spase" union type.
41pub type UnionArray = Array<Union>;
42
43/// The canonical Union encoding.
44///
45/// This is similar to Arrow's "spase" union type.
46#[derive(Clone, Debug)]
47pub struct Union;
48
49impl VTable for Union {
50    type TypedArrayData = EmptyArrayData;
51    type OperationsVTable = Self;
52    type ValidityVTable = Self;
53
54    fn id(&self) -> ArrayId {
55        static ID: CachedId = CachedId::new("vortex.union");
56        *ID
57    }
58
59    fn validate(
60        &self,
61        _data: &EmptyArrayData,
62        dtype: &DType,
63        len: usize,
64        slots: &[Option<ArrayRef>],
65    ) -> VortexResult<()> {
66        let type_ids = slots
67            .get(TYPE_IDS_SLOT)
68            .and_then(Option::as_ref)
69            .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?;
70        let variant_arrays = slots
71            .get(CHILDREN_OFFSET..)
72            .unwrap_or_default()
73            .iter()
74            .enumerate()
75            .map(|(index, slot)| {
76                slot.as_ref()
77                    .ok_or_else(|| vortex_err!("UnionArray is missing child {index}"))
78            })
79            .collect::<VortexResult<Vec<_>>>()?;
80
81        validate::validate_union_components(type_ids, &variant_arrays, dtype, len)
82    }
83
84    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
85        0
86    }
87
88    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
89        vortex_panic!("UnionArray buffer index {idx} out of bounds")
90    }
91
92    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
93        vortex_panic!("UnionArray buffer_name index {idx} out of bounds")
94    }
95
96    fn with_buffers(
97        &self,
98        array: ArrayView<'_, Self>,
99        buffers: &[BufferHandle],
100    ) -> VortexResult<ArrayParts<Self>> {
101        with_empty_buffers(self, array, buffers)
102    }
103
104    fn serialize(
105        _array: ArrayView<'_, Self>,
106        _session: &VortexSession,
107    ) -> VortexResult<Option<Vec<u8>>> {
108        Ok(Some(vec![]))
109    }
110
111    fn deserialize(
112        &self,
113        dtype: &DType,
114        len: usize,
115        metadata: &[u8],
116        buffers: &[BufferHandle],
117        children: &dyn ArrayChildren,
118        _session: &VortexSession,
119    ) -> VortexResult<ArrayParts<Self>> {
120        vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata");
121        vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers");
122        let DType::Union(variants, nullability) = dtype else {
123            vortex_bail!("Expected union dtype, found {dtype}")
124        };
125        vortex_ensure_eq!(
126            children.len(),
127            CHILDREN_OFFSET + variants.len(),
128            "UnionArray expected {} children, found {}",
129            CHILDREN_OFFSET + variants.len(),
130            children.len()
131        );
132
133        let type_ids = children.get(TYPE_IDS_SLOT, &union_type_ids_dtype(*nullability), len)?;
134        let sparse_children = variants
135            .variants()
136            .enumerate()
137            .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len))
138            .collect::<VortexResult<Vec<_>>>()?;
139
140        Ok(make_union_parts(
141            type_ids,
142            variants.clone(),
143            &sparse_children,
144        ))
145    }
146
147    fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String {
148        if idx == TYPE_IDS_SLOT {
149            "type_ids".to_string()
150        } else {
151            array.variants().names()[idx - CHILDREN_OFFSET].to_string()
152        }
153    }
154
155    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
156        Ok(ExecutionResult::done(array))
157    }
158
159    fn append_to_builder(
160        _array: ArrayView<'_, Self>,
161        _builder: &mut dyn ArrayBuilder,
162        _ctx: &mut ExecutionCtx,
163    ) -> VortexResult<()> {
164        todo!("TODO(connor)[Union]: implement append_to_builder for Union arrays")
165    }
166
167    fn reduce_parent(
168        array: ArrayView<'_, Self>,
169        parent: &ArrayRef,
170        child_idx: usize,
171    ) -> VortexResult<Option<ArrayRef>> {
172        PARENT_RULES.evaluate(array, parent, child_idx)
173    }
174}
vortex_array/arrays/union/vtable/mod.rs

vortex_array/arrays/union/vtable/
mod.rs
