Add UnionArray (#8834)

## Rationale for this change

Tracking issue: #8769

This is the array layer of the Union stack. The compute integration has
been split into the stacked follow-up #8884.

## What changes are included in this PR?

- Adds the canonical sparse `UnionArray` encoding with nullable u8 type
IDs and row-aligned children.
- Adds construction and invariant validation, typed accessors, scalar
access, top-level validity, serialization, and session registration.
- Registers UnionArray as the canonical representation for Union dtypes.
- Leaves filter, take, slice, mask, cast, aggregate, constant/chunked
canonicalization, and compression support to #8884.

I think that this is a mostly straightforward implementation. The one
implementation thing I will shout out is that we do not validate that
the entire type IDs array has valid type IDs. If we come across an
invalid type ID, we just panic, and that is similar behavior to
`DictArray` validation.

## What APIs are changed? Are there any user-facing changes?

Adds the public Union, UnionArray, UnionArrayExt, and UnionDataParts
APIs.

---------

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

Author: connortsui20

fuzz/src/array/fill_null.rs

@@ -49,6 +49,9 @@ pub fn fill_null_canonical_array(
         | Canonical::List(_)
         | Canonical::FixedSizeList(_)
         | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?,
+        Canonical::Union(_) => {
+            todo!("TODO(connor)[Union]: support Union arrays in the fill_null fuzzer")
+        }
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/mask.rs

@@ -149,6 +149,9 @@ pub fn mask_canonical_array(
                 .with_nullability(masked_storage.dtype().nullability());
             ExtensionArray::new(ext_dtype, masked_storage).into_array()
         }
+        Canonical::Union(_) => {
+            todo!("TODO(connor)[Union]: support Union arrays in the mask fuzzer")
+        }
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/scalar_at.rs

@@ -104,6 +104,9 @@ 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(_) => {
+            todo!("TODO(connor)[Union]: support Union arrays in the scalar_at fuzzer")
+        }
         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(_) => {
+                        todo!("TODO(connor)[Union]: implement IsConstant for Union arrays")
+                    }
                     Canonical::Variant(_) => {
                         vortex_bail!("Variant arrays don't support IsConstant")
                     }

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

@@ -417,6 +417,9 @@ impl AggregateFnVTable for MinMax {
                 Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx),
                 Canonical::Extension(e) => accumulate_extension(partial, e, ctx),
                 Canonical::Null(_) => Ok(()),
+                Canonical::Union(_) => {
+                    todo!("TODO(connor)[Union]: implement min_max for Union arrays")
+                }
                 Canonical::Struct(_)
                 | Canonical::List(_)
                 | Canonical::FixedSizeList(_)

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

@@ -199,6 +199,9 @@ 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::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx),
         Canonical::Variant(_) => {
             vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")

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

@@ -53,6 +53,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::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

@@ -95,6 +95,9 @@ 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::Extension(a) => {
             let filtered_storage = a
                 .storage_array()

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

@@ -50,6 +50,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(_) => {
+            todo!("TODO(connor)[Union]: implement masking for Union arrays")
+        }
         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
@@ -112,6 +112,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;