Add UnionArray

fuzz/src/array/fill_null.rs

@@ -49,6 +49,7 @@ pub fn fill_null_canonical_array(
         | Canonical::List(_)
         | Canonical::FixedSizeList(_)
         | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?,
+        Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/mask.rs

@@ -149,6 +149,7 @@ pub fn mask_canonical_array(
                 .with_nullability(masked_storage.dtype().nullability());
             ExtensionArray::new(ext_dtype, masked_storage).into_array()
         }
+        Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/scalar_at.rs

@@ -104,6 +104,7 @@ pub fn scalar_at_canonical_array(
             let storage_scalar = scalar_at_canonical_array(storage_canonical, index, ctx)?;
             Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar)
         }
+        Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

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

@@ -404,6 +404,9 @@ impl AggregateFnVTable for IsConstant {
                     Canonical::List(l) => check_listview_constant(l, ctx)?,
                     Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
                     Canonical::Null(_) => true,
+                    Canonical::Union(_) => {
+                        vortex_bail!("Union arrays don't support IsConstant yet")
+                    }
                     Canonical::Variant(_) => {
                         vortex_bail!("Variant arrays don't support IsConstant")
                     }

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

@@ -420,6 +420,7 @@ impl AggregateFnVTable for MinMax {
                 Canonical::Struct(_)
                 | Canonical::List(_)
                 | Canonical::FixedSizeList(_)
+                | Canonical::Union(_)
                 | Canonical::Variant(_) => {
                     vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype())
                 }

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

@@ -9,6 +9,7 @@ mod list_view;
 mod null;
 mod primitive;
 mod struct_;
+mod union;
 mod varbinview;
 
 use std::mem::size_of;
@@ -21,6 +22,7 @@ use list_view::list_view_uncompressed_size_in_bytes;
 use null::null_uncompressed_size_in_bytes;
 use primitive::primitive_uncompressed_size_in_bytes;
 use struct_::struct_uncompressed_size_in_bytes;
+use union::union_uncompressed_size_in_bytes;
 use varbinview::varbinview_uncompressed_size_in_bytes;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -199,6 +201,7 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
         Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
         Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
         Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx),
+        Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx),
         Canonical::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx),
         Canonical::Variant(_) => {
             vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
@@ -229,11 +232,14 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
             array.len(),
             array.scalar().as_binary().value().map(|value| value.len()),
         )?,
-        DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => {
+        DType::List(..)
+        | DType::FixedSizeList(..)
+        | DType::Struct(..)
+        | DType::Union(..)
+        | DType::Extension(_) => {
             let canonical = array.array().clone().execute::<Canonical>(ctx)?;
             return canonical_uncompressed_size_in_bytes(&canonical, ctx);
         }
-        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::Variant(_) => {
             vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
         }

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

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use super::uncompressed_size_in_bytes_u64;
+use crate::ExecutionCtx;
+use crate::arrays::UnionArray;
+use crate::arrays::union::UnionArrayExt;
+
+pub(super) fn union_uncompressed_size_in_bytes(
+    array: &UnionArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let mut size = uncompressed_size_in_bytes_u64(array.type_ids(), ctx)?;
+
+    for child in array.iter_children() {
+        size = size
+            .checked_add(uncompressed_size_in_bytes_u64(child, ctx)?)
+            .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+    }
+
+    Ok(size)
+}

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

@@ -19,11 +19,14 @@ use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListViewArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::StructArray;
+use crate::arrays::UnionArray;
 use crate::arrays::VariantArray;
 use crate::arrays::chunked::ChunkedArrayExt;
 use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::listview::ListViewRebuildMode;
+use crate::arrays::union::TYPE_IDS_DTYPE;
+use crate::arrays::union::UnionArrayExt;
 use crate::arrays::variant::VariantArrayExt;
 use crate::builders::builder_with_capacity_in;
 use crate::builtins::ArrayBuiltins;
@@ -54,6 +57,7 @@ pub(super) fn _canonicalize(
             let struct_array = pack_struct_chunks(owned_chunks, ctx)?;
             Canonical::Struct(struct_array)
         }
+        DType::Union(..) => Canonical::Union(pack_union_chunks(owned_chunks, ctx)?),
         DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks(
             &owned_chunks,
             array.array().validity()?,
@@ -89,6 +93,46 @@ fn pack_struct_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexRe
         .process_results(|iter| StructArray::try_concat(iter))?
 }
 
+/// Packs sparse union chunks into one union with chunked type IDs and sparse children.
+fn pack_union_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexResult<UnionArray> {
+    let union_chunks = chunks
+        .into_iter()
+        .map(|chunk| chunk.execute::<UnionArray>(ctx))
+        .collect::<VortexResult<Vec<_>>>()?;
+    let variants = union_chunks[0].variants().clone();
+    let type_ids = ChunkedArray::try_new(
+        union_chunks
+            .iter()
+            .map(|chunk| chunk.type_ids().clone())
+            .collect(),
+        TYPE_IDS_DTYPE,
+    )?
+    .into_array();
+    let children = variants
+        .variants()
+        .enumerate()
+        .map(|(index, dtype)| {
+            ChunkedArray::try_new(
+                union_chunks
+                    .iter()
+                    .map(|chunk| {
+                        chunk
+                            .child(index)
+                            .vortex_expect("variant index must have a sparse child")
+                            .clone()
+                    })
+                    .collect(),
+                dtype,
+            )
+            .map(IntoArray::into_array)
+        })
+        .collect::<VortexResult<Vec<_>>>()?;
+
+    // SAFETY: Every component was taken from validated UnionArrays with the same dtype and
+    // chunked along identical row boundaries.
+    Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) })
+}
+
 /// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children.
 ///
 /// The caller guarantees there are at least 2 chunks.

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

@@ -251,9 +251,12 @@ impl VTable for Chunked {
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         match array.dtype() {
-            // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path
-            // cannot express.
-            DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => {
+            // Nested canonical arrays need child swizzling that the builder path cannot express.
+            DType::Struct(..)
+            | DType::List(..)
+            | DType::FixedSizeList(..)
+            | DType::Union(..)
+            | DType::Variant(..) => {
                 // TODO(joe)[#7674]: iterative execution here too
                 Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?))
             }

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

@@ -8,6 +8,7 @@ use vortex_buffer::Buffer;
 use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -23,13 +24,15 @@ use crate::arrays::ListViewArray;
 use crate::arrays::NullArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::StructArray;
+use crate::arrays::UnionArray;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::VariantArray;
 use crate::arrays::varbinview::BinaryView;
 use crate::builders::builder_with_capacity;
 use crate::dtype::DType;
 use crate::dtype::DecimalType;
 use crate::dtype::Nullability;
+use crate::dtype::UnionVariants;
 use crate::match_each_decimal_value;
 use crate::match_each_decimal_value_type;
 use crate::match_each_native_ptype;
@@ -164,7 +167,9 @@ pub(crate) fn constant_canonicalize(
                 StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
             })
         }
-        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+        DType::Union(variants) => {
+            Canonical::Union(constant_canonical_union(scalar, variants, array.len())?)
+        }
         DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
             array.array().clone().into_array(),
             None,
@@ -188,6 +193,51 @@ pub(crate) fn constant_canonicalize(
     })
 }
 
+fn constant_canonical_union(
+    scalar: &Scalar,
+    variants: &UnionVariants,
+    len: usize,
+) -> VortexResult<UnionArray> {
+    if scalar.is_null() {
+        vortex_bail!(
+            "Canonicalizing a null union scalar is not supported until null semantics are defined"
+        )
+    }
+    if variants.variants().any(|variant| variant.is_nullable()) {
+        vortex_bail!("Canonical UnionArray children must be non-nullable")
+    }
+
+    let union = scalar.as_union();
+    let type_id = union
+        .type_id()
+        .vortex_expect("non-null union scalar must have a type ID");
+    let selected_child = union
+        .child_index()
+        .vortex_expect("validated union scalar must select a child");
+    let selected_value = union
+        .value()
+        .vortex_expect("non-null union scalar must have a value");
+
+    let children = variants
+        .variants()
+        .enumerate()
+        .map(|(index, dtype)| {
+            let value = if index == selected_child {
+                selected_value.clone()
+            } else {
+                Scalar::zero_value(&dtype)
+            };
+            ConstantArray::new(value, len).into_array()
+        })
+        .collect::<Vec<_>>();
+
+    UnionArray::try_new(
+        ConstantArray::new(type_id, len).into_array(),
+        variants.clone(),
+        children,
+    )
+}
+
 fn constant_canonical_byte_view(
     scalar_bytes: Option<&[u8]>,
     dtype: &DType,

vortex-array/src/arrays/dict/execute.rs

@@ -30,6 +30,7 @@ use crate::arrays::VarBinViewArray;
 use crate::arrays::VariantArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::dict::TakeReduce;
+use crate::arrays::union::compute::take_union;
 use crate::arrays::variant::VariantArrayExt;
 
 /// Take from a canonical array using indices (codes), returning a new canonical array.
@@ -53,6 +54,10 @@ pub(crate) fn take_canonical(
             Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx))
         }
         Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
+        Canonical::Union(a) => {
+            let indices = codes.clone().into_array();
+            Canonical::Union(take_union(a.as_view(), &indices)?)
+        }
         Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
         Canonical::Variant(a) => {
             let indices = codes.clone().into_array();

vortex-array/src/arrays/filter/execute/mod.rs

@@ -24,6 +24,7 @@ use crate::arrays::NullArray;
 use crate::arrays::VariantArray;
 use crate::arrays::extension::ExtensionArrayExt;
 use crate::arrays::filter::FilterArrayExt;
+use crate::arrays::union::compute::filter_union;
 use crate::arrays::variant::VariantArrayExt;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
@@ -95,6 +96,13 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
             Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask))
         }
         Canonical::Struct(a) => Canonical::Struct(struct_::filter_struct(&a, mask)),
+        Canonical::Union(a) => {
+            let mask = Mask::Values(Arc::clone(mask));
+            Canonical::Union(
+                filter_union(a.as_view(), &mask)
+                    .vortex_expect("UnionArray children must support filter"),
+            )
+        }
         Canonical::Extension(a) => {
             let filtered_storage = a
                 .storage_array()

vortex-array/src/arrays/masked/execute.rs

@@ -6,6 +6,7 @@
 use std::sync::Arc;
 
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::IntoArray;
@@ -50,6 +51,9 @@ pub fn mask_validity_canonical(
             Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?)
         }
         Canonical::Struct(a) => Canonical::Struct(mask_validity_struct(a, validity)?),
+        Canonical::Union(_) => {
+            vortex_bail!("Masking UnionArray is not supported until null semantics are defined")
+        }
         Canonical::Extension(a) => Canonical::Extension(mask_validity_extension(a, validity, ctx)?),
         Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity, ctx)?),
     })

vortex-array/src/arrays/mod.rs

@@ -5,8 +5,8 @@
 //!
 //! Canonical arrays are the default uncompressed representation for a logical dtype:
 //! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], [`VarBinViewArray`],
-//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`ExtensionArray`], and
-//! [`VariantArray`].
+//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`UnionArray`],
+//! [`ExtensionArray`], and [`VariantArray`].
 //!
 //! Utility and lazy arrays represent common transformations without immediately materializing
 //! their result. Examples include [`ChunkedArray`] for concatenation, [`ConstantArray`] for repeated
@@ -108,6 +108,10 @@ pub mod struct_;
 pub use struct_::Struct;
 pub use struct_::StructArray;
 
+pub mod union;
+pub use union::Union;
+pub use union::UnionArray;
+
 pub mod varbin;
 pub use varbin::VarBin;
 pub use varbin::VarBinArray;