Add UnionArray

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;