Add UnionArray

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;

vortex-array/src/arrays/union/array.rs

@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::iter::once;
+use std::sync::Arc;
+
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
+
+use crate::ArrayRef;
+use crate::ArraySlots;
+use crate::IntoArray;
+use crate::array::Array;
+use crate::array::ArrayParts;
+use crate::array::EmptyArrayData;
+use crate::array::TypedArrayRef;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::Union;
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::UnionVariants;
+
+/// The row-aligned array of type IDs selecting a union child.
+pub(super) const TYPE_IDS_SLOT: usize = 0;
+/// The offset at which the sparse child arrays begin in the slots vector.
+pub(super) const CHILDREN_OFFSET: usize = 1;
+
+pub(super) fn make_union_parts(
+    type_ids: ArrayRef,
+    variants: UnionVariants,
+    children: &[ArrayRef],
+) -> ArrayParts<Union> {
+    let len = type_ids.len();
+    let nullability = type_ids.dtype().nullability();
+    let slots: ArraySlots = once(Some(type_ids))
+        .chain(children.iter().cloned().map(Some))
+        .collect();
+
+    ArrayParts::new(
+        Union,
+        DType::Union(variants, nullability),
+        len,
+        EmptyArrayData,
+    )
+    .with_slots(slots)
+}
+
+/// Concrete parts of a [`UnionArray`](super::UnionArray).
+pub struct UnionDataParts {
+    /// The union variant schema.
+    pub variants: UnionVariants,
+    /// The row-aligned type IDs.
+    pub type_ids: ArrayRef,
+    /// The row-aligned sparse children in variant order.
+    pub children: Arc<[ArrayRef]>,
+}
+
+/// Accessors for a canonical sparse union array.
+pub trait UnionArrayExt: TypedArrayRef<Union> {
+    /// The union's variant schema.
+    fn variants(&self) -> &UnionVariants {
+        match self.as_ref().dtype() {
+            DType::Union(variants, _) => variants,
+            _ => unreachable!("UnionArrayExt requires a union dtype"),
+        }
+    }
+
+    /// The row-aligned `u8` type IDs whose nulls represent outer union nulls.
+    fn type_ids(&self) -> &ArrayRef {
+        self.as_ref().slots()[TYPE_IDS_SLOT]
+            .as_ref()
+            .vortex_expect("UnionArray type_ids slot")
+    }
+
+    /// Iterate over sparse children in variant order.
+    fn iter_children(&self) -> impl ExactSizeIterator<Item = &ArrayRef> + '_ {
+        self.as_ref().slots()[CHILDREN_OFFSET..]
+            .iter()
+            .map(|slot| slot.as_ref().vortex_expect("UnionArray child slot"))
+    }
+
+    /// Return the sparse children in variant order.
+    fn children(&self) -> Arc<[ArrayRef]> {
+        self.iter_children().cloned().collect()
+    }
+
+    /// Return a sparse child by its variant index.
+    fn child(&self, index: usize) -> Option<&ArrayRef> {
+        self.as_ref().slots().get(CHILDREN_OFFSET + index)?.as_ref()
+    }
+
+    /// Return a sparse child selected by a data-level type ID.
+    fn child_by_type_id(&self, type_id: u8) -> Option<&ArrayRef> {
+        self.child(self.variants().tag_to_child_index(type_id)?)
+    }
+
+    /// Return a sparse child selected by its variant name, if present.
+    fn child_by_name_opt(&self, name: impl AsRef<str>) -> Option<&ArrayRef> {
+        self.child(self.variants().find(name)?)
+    }
+
+    /// Return a sparse child selected by its variant name.
+    fn child_by_name(&self, name: impl AsRef<str>) -> VortexResult<&ArrayRef> {
+        let name = name.as_ref();
+        self.child_by_name_opt(name).ok_or_else(|| {
+            vortex_err!(
+                "Variant {name} not found in union array with names {:?}",
+                self.variants().names()
+            )
+        })
+    }
+}
+impl<T: TypedArrayRef<Union>> UnionArrayExt for T {}
+
+impl Array<Union> {
+    /// Construct a canonical sparse union array.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the components do not satisfy the invariants documented on
+    /// [`Self::new_unchecked`].
+    pub fn new(
+        type_ids: ArrayRef,
+        variants: UnionVariants,
+        children: impl Into<Arc<[ArrayRef]>>,
+    ) -> Self {
+        Self::try_new(type_ids, variants, children).vortex_expect("UnionArray construction failed")
+    }
+
+    /// Try to construct a canonical sparse union array.
+    ///
+    /// Type ID values are not validated during construction. Accessing a non-null row whose type
+    /// ID is not declared by `variants` will panic.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if `type_ids` is not a `u8` array, or if the sparse children do not match
+    /// the variant schema and outer array length.
+    pub fn try_new(
+        type_ids: ArrayRef,
+        variants: UnionVariants,
+        children: impl Into<Arc<[ArrayRef]>>,
+    ) -> VortexResult<Self> {
+        vortex_ensure!(
+            matches!(
+                type_ids.dtype(),
+                DType::Primitive(crate::dtype::PType::U8, _)
+            ),
+            "UnionArray type_ids must be u8, got {}",
+            type_ids.dtype()
+        );
+
+        let children = children.into();
+        Array::try_from_parts(make_union_parts(type_ids, variants, &children))
+    }
+
+    /// Construct a canonical sparse union array without validation.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure `type_ids` is a `u8` array, every child has the corresponding variant
+    /// dtype, and all arrays have the same length. Null type IDs represent outer union nulls.
+    pub unsafe fn new_unchecked(
+        type_ids: ArrayRef,
+        variants: UnionVariants,
+        children: impl Into<Arc<[ArrayRef]>>,
+    ) -> Self {
+        let children = children.into();
+        unsafe { Array::from_parts_unchecked(make_union_parts(type_ids, variants, &children)) }
+    }
+
+    /// Deconstruct this array into its type IDs, variant schema, and sparse children.
+    pub fn into_data_parts(self) -> UnionDataParts {
+        let variants = self.variants().clone();
+        let type_ids = self.type_ids().clone();
+        let children = self.children();
+        UnionDataParts {
+            variants,
+            type_ids,
+            children,
+        }
+    }
+
+    /// Create an empty array for a union dtype.
+    pub(crate) fn empty(variants: UnionVariants, nullability: Nullability) -> Self {
+        let type_ids = PrimitiveArray::empty::<u8>(nullability).into_array();
+        let children: Vec<_> = variants
+            .variants()
+            .map(|dtype| crate::Canonical::empty(&dtype).into_array())
+            .collect();
+
+        Self::new(type_ids, variants, children)
+    }
+}

vortex-array/src/arrays/union/mod.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Canonical sparse union arrays.
+//!
+//! A [`UnionArray`] stores one `u8` type ID per row followed by one row-aligned child for each
+//! variant. The type ID selects which child's value is active for a row; values in all other
+//! children at that row are placeholders. Outer union nulls are stored as nulls in the type IDs
+//! child, independently of nulls in the selected variant child.
+//!
+//! Type ID values are not validated during construction. Accessing a non-null row whose type ID
+//! is not declared by the union variants will panic.
+
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::PType;
+
+mod array;
+pub use array::UnionArrayExt;
+pub use array::UnionDataParts;
+pub use vtable::UnionArray;
+
+mod vtable;
+pub use vtable::Union;
+
+pub(crate) fn union_type_ids_dtype(nullability: Nullability) -> DType {
+    DType::Primitive(PType::U8, nullability)
+}
+
+#[cfg(test)]
+mod tests;