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

vortex_array/arrays/union/vtable/
mod.rs
