Add nullable 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::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 = Scalar::try_default_value(&dt).ok_or_else(|| {
+                                vortex_err!(
+                                    "cannot canonicalize null constant struct: field dtype {dt} \
+                                     has no default 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,51 @@ 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.value());
+
+    let children = variants
+        .variants()
+        .enumerate()
+        .map(|(index, dtype)| {
+            let value = match selected {
+                Some((selected_child, selected_value)) if index == selected_child => {
+                    selected_value.clone()
+                }
+                _ => Scalar::try_default_value(&dtype).ok_or_else(|| {
+                    vortex_err!(
+                        "cannot canonicalize constant union: unselected variant {} ({dtype}) has \
+                         no default 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,
+    )
+}
+
 fn constant_canonical_byte_view(
     scalar_bytes: Option<&[u8]>,
     dtype: &DType,

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

@@ -18,3 +18,47 @@ impl OperationsVTable<Constant> for Constant {
         Ok(array.scalar.clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+
+    use crate::IntoArray;
+    use crate::VortexSessionExecute;
+    use crate::arrays::ConstantArray;
+    use crate::dtype::DType;
+    use crate::dtype::Nullability;
+    use crate::dtype::PType;
+    use crate::dtype::UnionVariants;
+    use crate::scalar::Scalar;
+
+    #[test]
+    fn scalar_at_preserves_union_scalar() -> VortexResult<()> {
+        let variants = UnionVariants::try_new(
+            ["int", "string"].into(),
+            vec![
+                DType::Primitive(PType::I32, Nullability::Nullable),
+                DType::Utf8(Nullability::NonNullable),
+            ],
+            vec![5, 9],
+        )?;
+
+        let scalar = Scalar::union(
+            variants.clone(),
+            5,
+            Scalar::primitive(42_i32, Nullability::Nullable),
+            Nullability::Nullable,
+        )?;
+        let array = ConstantArray::new(scalar.clone(), 3).into_array();
+        let mut ctx = crate::array_session().create_execution_ctx();
+
+        assert_eq!(array.execute_scalar(1, &mut ctx)?, scalar);
+
+        let null = Scalar::null(DType::Union(variants, Nullability::Nullable));
+        let array = ConstantArray::new(null.clone(), 3).into_array();
+
+        assert_eq!(array.execute_scalar(1, &mut ctx)?, null);
+
+        Ok(())
+    }
+}

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();