Add UnionArray compute functions

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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,9 +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(_) => {
-            todo!("TODO(connor)[Union]: implement UncompressedSizeInBytes for Union arrays")
-        }
+        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")
@@ -232,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::UnionArrayExt;
+use crate::arrays::union::type_ids_dtype;
 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(union_chunks[0].dtype().nullability()),
+    )?
+    .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_err;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -23,13 +24,16 @@ 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::PType;
+use crate::dtype::UnionVariants;
 use crate::match_each_decimal_value;
 use crate::match_each_decimal_value_type;
 use crate::match_each_native_ptype;
@@ -146,16 +150,22 @@ pub(crate) fn constant_canonicalize(
                     .collect(),
                 None => {
                     assert!(matches!(validity, Validity::AllInvalid));
-                    // The struct is entirely null, so fields just need placeholder values with the
-                    // correct dtype. We use `default_value` which returns a zero for non-nullable
-                    // dtypes and null for nullable dtypes, preserving each field's nullability.
                     struct_dtype
                         .fields()
                         .map(|dt| {
-                            let scalar = Scalar::default_value(&dt);
-                            ConstantArray::new(scalar, array.len()).into_array()
+                            if array.is_empty() {
+                                return Ok(Canonical::empty(&dt).into_array());
+                            }
+
+                            let scalar = try_placeholder_scalar(&dt).ok_or_else(|| {
+                                vortex_err!(
+                                    "cannot canonicalize null constant struct: field dtype {dt} \
+                                     has no placeholder value"
+                                )
+                            })?;
+                            Ok(ConstantArray::new(scalar, array.len()).into_array())
                         })
-                        .collect()
+                        .collect::<VortexResult<Vec<_>>>()?
                 }
             };
             // SAFETY: Fields are constructed from the same struct scalar, all have same
@@ -164,7 +174,12 @@ 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, nullability) => Canonical::Union(constant_canonical_union(
+            scalar,
+            variants,
+            *nullability,
+            array.len(),
+        )?),
         DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
             array.array().clone().into_array(),
             None,
@@ -188,6 +203,97 @@ pub(crate) fn constant_canonicalize(
     })
 }
 
+fn constant_canonical_union(
+    scalar: &Scalar,
+    variants: &UnionVariants,
+    nullability: Nullability,
+    len: usize,
+) -> VortexResult<UnionArray> {
+    if len == 0 {
+        return Ok(UnionArray::empty(variants.clone(), nullability));
+    }
+
+    let union = scalar.as_union();
+    let selected = union.child_index().zip(union.child());
+
+    let children = variants
+        .variants()
+        .enumerate()
+        .map(|(index, dtype)| {
+            let value = match &selected {
+                Some((selected_child, selected_value)) if index == *selected_child => {
+                    selected_value.clone()
+                }
+                _ => try_placeholder_scalar(&dtype).ok_or_else(|| {
+                    vortex_err!(
+                        "cannot canonicalize constant union: unselected variant {} ({dtype}) has \
+                         no placeholder value",
+                        variants.names()[index]
+                    )
+                })?,
+            };
+            Ok(ConstantArray::new(value, len).into_array())
+        })
+        .collect::<VortexResult<Vec<_>>>()?;
+
+    let type_id = match union.type_id() {
+        Some(type_id) => Scalar::primitive(type_id, nullability),
+        None => Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
+    };
+
+    UnionArray::try_new(
+        ConstantArray::new(type_id, len).into_array(),
+        variants.clone(),
+        children,
+    )
+}
+
+/// Builds an arbitrary valid scalar for array positions ignored by a parent validity or type ID.
+/// This does not give the dtype a semantic default value.
+fn try_placeholder_scalar(dtype: &DType) -> Option<Scalar> {
+    if let Some(default) = Scalar::try_default_value(dtype) {
+        return Some(default);
+    }
+
+    match dtype {
+        DType::FixedSizeList(element_dtype, size, nullability) => {
+            let element = try_placeholder_scalar(element_dtype)?;
+            Some(Scalar::fixed_size_list(
+                Arc::clone(element_dtype),
+                vec![element; *size as usize],
+                *nullability,
+            ))
+        }
+        DType::Struct(fields, _) => {
+            let children = fields
+                .fields()
+                .map(|field| try_placeholder_scalar(&field))
+                .collect::<Option<Vec<_>>>()?;
+            Some(Scalar::struct_(dtype.clone(), children))
+        }
+        DType::Union(variants, nullability) => {
+            variants
+                .variants()
+                .enumerate()
+                .find_map(|(index, child_dtype)| {
+                    let child = try_placeholder_scalar(&child_dtype)?;
+                    Scalar::union(
+                        variants.clone(),
+                        variants.child_index_to_tag(index),
+                        child,
+                        *nullability,
+                    )
+                    .ok()
+                })
+        }
+        DType::Extension(ext_dtype) => {
+            let storage = try_placeholder_scalar(ext_dtype.storage_dtype())?;
+            Some(Scalar::extension_ref(ext_dtype.clone(), storage))
+        }
+        _ => None,
+    }
+}
+
 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,8 +54,9 @@ 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(_) => {
-            todo!("TODO(connor)[Union]: implement dictionary execution for Union arrays")
+        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) => {

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,8 +96,12 @@ 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(_) => {
-            todo!("TODO(connor)[Union]: implement filter for Union arrays")
+        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