Skip to content

Commit 21e3100

Browse files
connortsui20connortsui20
committed
Add UnionArray
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 5e0113d commit 21e3100

32 files changed

Lines changed: 1053 additions & 14 deletions

File tree

‎fuzz/src/array/fill_null.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ pub fn fill_null_canonical_array(
4949
| Canonical::List(_)
5050
| Canonical::FixedSizeList(_)
5151
| Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?,
52+
Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
5253
Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
5354
})
5455
}

‎fuzz/src/array/mask.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ pub fn mask_canonical_array(
149149
.with_nullability(masked_storage.dtype().nullability());
150150
ExtensionArray::new(ext_dtype, masked_storage).into_array()
151151
}
152+
Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
152153
Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
153154
})
154155
}

‎fuzz/src/array/scalar_at.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ pub fn scalar_at_canonical_array(
104104
let storage_scalar = scalar_at_canonical_array(storage_canonical, index, ctx)?;
105105
Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar)
106106
}
107+
Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
107108
Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
108109
})
109110
}

‎vortex-array/src/aggregate_fn/fns/is_constant/mod.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,9 @@ impl AggregateFnVTable for IsConstant {
404404
Canonical::List(l) => check_listview_constant(l, ctx)?,
405405
Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
406406
Canonical::Null(_) => true,
407+
Canonical::Union(_) => {
408+
vortex_bail!("Union arrays don't support IsConstant yet")
409+
}
407410
Canonical::Variant(_) => {
408411
vortex_bail!("Variant arrays don't support IsConstant")
409412
}

‎vortex-array/src/aggregate_fn/fns/min_max/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,7 @@ impl AggregateFnVTable for MinMax {
420420
Canonical::Struct(_)
421421
| Canonical::List(_)
422422
| Canonical::FixedSizeList(_)
423+
| Canonical::Union(_)
423424
| Canonical::Variant(_) => {
424425
vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype())
425426
}

‎vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ mod list_view;
99
mod null;
1010
mod primitive;
1111
mod struct_;
12+
mod union;
1213
mod varbinview;
1314

1415
use std::mem::size_of;
@@ -21,6 +22,7 @@ use list_view::list_view_uncompressed_size_in_bytes;
2122
use null::null_uncompressed_size_in_bytes;
2223
use primitive::primitive_uncompressed_size_in_bytes;
2324
use struct_::struct_uncompressed_size_in_bytes;
25+
use union::union_uncompressed_size_in_bytes;
2426
use varbinview::varbinview_uncompressed_size_in_bytes;
2527
use vortex_error::VortexExpect;
2628
use vortex_error::VortexResult;
@@ -199,6 +201,7 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
199201
Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
200202
Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
201203
Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx),
204+
Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx),
202205
Canonical::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx),
203206
Canonical::Variant(_) => {
204207
vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
@@ -229,11 +232,14 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
229232
array.len(),
230233
array.scalar().as_binary().value().map(|value| value.len()),
231234
)?,
232-
DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => {
235+
DType::List(..)
236+
| DType::FixedSizeList(..)
237+
| DType::Struct(..)
238+
| DType::Union(..)
239+
| DType::Extension(_) => {
233240
let canonical = array.array().clone().execute::<Canonical>(ctx)?;
234241
return canonical_uncompressed_size_in_bytes(&canonical, ctx);
235242
}
236-
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
237243
DType::Variant(_) => {
238244
vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
239245
}

‎vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexResult;
5+
use vortex_error::vortex_err;
6+
7+
use super::uncompressed_size_in_bytes_u64;
8+
use crate::ExecutionCtx;
9+
use crate::arrays::UnionArray;
10+
use crate::arrays::union::UnionArrayExt;
11+
12+
pub(super) fn union_uncompressed_size_in_bytes(
13+
array: &UnionArray,
14+
ctx: &mut ExecutionCtx,
15+
) -> VortexResult<u64> {
16+
let mut size = uncompressed_size_in_bytes_u64(array.type_ids(), ctx)?;
17+
18+
for child in array.iter_children() {
19+
size = size
20+
.checked_add(uncompressed_size_in_bytes_u64(child, ctx)?)
21+
.ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
22+
}
23+
24+
Ok(size)
25+
}

‎vortex-array/src/arrays/chunked/vtable/canonical.rs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,13 @@ use crate::arrays::FixedSizeListArray;
1919
use crate::arrays::ListViewArray;
2020
use crate::arrays::PrimitiveArray;
2121
use crate::arrays::StructArray;
22+
use crate::arrays::UnionArray;
2223
use crate::arrays::VariantArray;
2324
use crate::arrays::chunked::ChunkedArrayExt;
2425
use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
2526
use crate::arrays::listview::ListViewArrayExt;
2627
use crate::arrays::listview::ListViewRebuildMode;
28+
use crate::arrays::union::UnionArrayExt;
2729
use crate::arrays::variant::VariantArrayExt;
2830
use crate::builders::builder_with_capacity_in;
2931
use crate::builtins::ArrayBuiltins;
@@ -54,6 +56,7 @@ pub(super) fn _canonicalize(
5456
let struct_array = pack_struct_chunks(owned_chunks, ctx)?;
5557
Canonical::Struct(struct_array)
5658
}
59+
DType::Union(..) => Canonical::Union(pack_union_chunks(owned_chunks, ctx)?),
5760
DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks(
5861
&owned_chunks,
5962
array.array().validity()?,
@@ -89,6 +92,47 @@ fn pack_struct_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexRe
8992
.process_results(|iter| StructArray::try_concat(iter))?
9093
}
9194

95+
/// Packs sparse union chunks into one union with chunked type IDs and sparse children.
96+
fn pack_union_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexResult<UnionArray> {
97+
let union_chunks = chunks
98+
.into_iter()
99+
.map(|chunk| chunk.execute::<UnionArray>(ctx))
100+
.collect::<VortexResult<Vec<_>>>()?;
101+
let variants = union_chunks[0].variants().clone();
102+
let type_ids = ChunkedArray::try_new(
103+
union_chunks
104+
.iter()
105+
.map(|chunk| chunk.type_ids().clone())
106+
.collect(),
107+
DType::Primitive(PType::I8, Nullability::NonNullable),
108+
)?
109+
.into_array();
110+
let children = (0..variants.len())
111+
.map(|index| {
112+
let dtype = variants
113+
.variant_by_index(index)
114+
.vortex_expect("variant index must have a dtype");
115+
ChunkedArray::try_new(
116+
union_chunks
117+
.iter()
118+
.map(|chunk| {
119+
chunk
120+
.child(index)
121+
.vortex_expect("variant index must have a sparse child")
122+
.clone()
123+
})
124+
.collect(),
125+
dtype,
126+
)
127+
.map(IntoArray::into_array)
128+
})
129+
.collect::<VortexResult<Vec<_>>>()?;
130+
131+
// SAFETY: Every component was taken from validated UnionArrays with the same dtype and
132+
// chunked along identical row boundaries.
133+
Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) })
134+
}
135+
92136
/// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children.
93137
///
94138
/// The caller guarantees there are at least 2 chunks.

‎vortex-array/src/arrays/chunked/vtable/mod.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -251,9 +251,12 @@ impl VTable for Chunked {
251251

252252
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
253253
match array.dtype() {
254-
// Struct, List, FixedSizeList, and Variant need child swizzling that the builder path
255-
// cannot express.
256-
DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => {
254+
// Nested canonical arrays need child swizzling that the builder path cannot express.
255+
DType::Struct(..)
256+
| DType::List(..)
257+
| DType::FixedSizeList(..)
258+
| DType::Union(..)
259+
| DType::Variant(..) => {
257260
// TODO(joe)[#7674]: iterative execution here too
258261
Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?))
259262
}

‎vortex-array/src/arrays/constant/vtable/canonical.rs‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use vortex_buffer::Buffer;
88
use vortex_buffer::buffer;
99
use vortex_error::VortexExpect;
1010
use vortex_error::VortexResult;
11+
use vortex_error::vortex_bail;
1112

1213
use crate::Canonical;
1314
use crate::ExecutionCtx;
@@ -23,6 +24,7 @@ use crate::arrays::ListViewArray;
2324
use crate::arrays::NullArray;
2425
use crate::arrays::PrimitiveArray;
2526
use crate::arrays::StructArray;
27+
use crate::arrays::UnionArray;
2628
use crate::arrays::VarBinViewArray;
2729
use crate::arrays::VariantArray;
2830
use crate::arrays::varbinview::BinaryView;
@@ -164,7 +166,49 @@ pub(crate) fn constant_canonicalize(
164166
StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
165167
})
166168
}
167-
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
169+
DType::Union(variants) => {
170+
if scalar.is_null() {
171+
vortex_bail!(
172+
"Canonicalizing a null union scalar is not supported until null semantics are defined"
173+
)
174+
}
175+
if variants.variants().any(|variant| variant.is_nullable()) {
176+
vortex_bail!("Canonical UnionArray children must be non-nullable")
177+
}
178+
179+
let union = scalar.as_union();
180+
let type_id = union
181+
.type_id()
182+
.vortex_expect("non-null union scalar must have a type ID");
183+
let child_index = union
184+
.child_index()
185+
.vortex_expect("validated union scalar must select a child");
186+
let selected_value = union
187+
.value()
188+
.vortex_expect("non-null union scalar must have a value");
189+
let children = variants
190+
.variants()
191+
.enumerate()
192+
.map(|(index, dtype)| {
193+
let value = if index == child_index {
194+
selected_value.clone()
195+
} else {
196+
Scalar::zero_value(&dtype)
197+
};
198+
ConstantArray::new(value, array.len()).into_array()
199+
})
200+
.collect::<Vec<_>>();
201+
202+
// SAFETY: The scalar's validated type ID selects `child_index`; all sparse children
203+
// have the declared dtype and the same length.
204+
Canonical::Union(unsafe {
205+
UnionArray::new_unchecked(
206+
ConstantArray::new(type_id, array.len()).into_array(),
207+
variants.clone(),
208+
children,
209+
)
210+
})
211+
}
168212
DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
169213
array.array().clone().into_array(),
170214
None,

0 commit comments

Comments
 (0)