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::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(TYPE_IDS_SLOT)
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(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!(
125            children.len(),
126            CHILDREN_OFFSET + variants.len(),
127            "UnionArray expected {} children, found {}",
128            CHILDREN_OFFSET + variants.len(),
129            children.len()
130        );
131
132        let type_ids = children.get(TYPE_IDS_SLOT, &union_type_ids_dtype(*nullability), len)?;
133        let sparse_children = variants
134            .variants()
135            .enumerate()
136            .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len))
137            .collect::<VortexResult<Vec<_>>>()?;
138
139        Ok(make_union_parts(
140            type_ids,
141            variants.clone(),
142            &sparse_children,
143        ))
144    }
145
146    fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String {
147        if idx == TYPE_IDS_SLOT {
148            "type_ids".to_string()
149        } else {
150            array.variants().names()[idx - CHILDREN_OFFSET].to_string()
151        }
152    }
153
154    fn execute(_array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
155        todo!("TODO(connor)[Union]: implement execute for Union arrays")
156    }
157
158    fn append_to_builder(
159        _array: ArrayView<'_, Self>,
160        _builder: &mut dyn ArrayBuilder,
161        _ctx: &mut ExecutionCtx,
162    ) -> VortexResult<()> {
163        todo!("TODO(connor)[Union]: implement append_to_builder for Union arrays")
164    }
165}
vortex_array/arrays/union/vtable/mod.rs

vortex_array/arrays/union/vtable/
mod.rs
