Add DType::Union and UnionVariants

encodings/sparse/src/canonical.rs

@@ -122,6 +122,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
             len,
             ctx,
         )?,
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::Decimal(decimal_dtype, nullability) => {
             let canonical_decimal_value_type =
                 DecimalType::smallest_decimal_value_type(decimal_dtype);

fuzz/src/array/compare.rs

@@ -173,7 +173,7 @@ pub fn compare_canonical_array(
             }))
             .into_array()
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/filter.rs

@@ -122,7 +122,7 @@ pub fn filter_canonical_array(
             }
             take_canonical_array_non_nullable_indices(array, indices.as_slice(), ctx)
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/mod.rs

@@ -474,6 +474,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
                     acc.intersection(&actions).copied().collect()
                 })
         }
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::List(..) | DType::FixedSizeList(..) => {
             // List supports: Compress, Slice, Take, Filter, MinMax, Mask, ScalarAt
             // Does NOT support: SearchSorted, Compare, Cast, Sum, FillNull

fuzz/src/array/search_sorted.rs

@@ -148,7 +148,7 @@ pub fn search_sorted_canonical_array(
                 .collect::<VortexResult<Vec<_>>>()?;
             scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side)
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/slice.rs

@@ -123,7 +123,7 @@ pub fn slice_canonical_array(
                 .into_array(),
             )
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/sort.rs

@@ -101,7 +101,7 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR
             });
             take_canonical_array_non_nullable_indices(array, &sort_indices, ctx)
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/take.rs

@@ -147,7 +147,7 @@ pub fn take_canonical_array(
             }
             Ok(builder.finish())
         }
-        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
+        d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

vortex-array/public-api.lock

@@ -9346,6 +9346,8 @@ pub vortex_array::dtype::DType::Primitive(vortex_array::dtype::PType, vortex_arr
 
 pub vortex_array::dtype::DType::Struct(vortex_array::dtype::StructFields, vortex_array::dtype::Nullability)
 
+pub vortex_array::dtype::DType::Union(vortex_array::dtype::UnionVariants, vortex_array::dtype::Nullability)
+
 pub vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability)
 
 pub vortex_array::dtype::DType::Variant(vortex_array::dtype::Nullability)
@@ -9376,6 +9378,8 @@ pub fn vortex_array::dtype::DType::as_struct_fields(&self) -> &vortex_array::dty
 
 pub fn vortex_array::dtype::DType::as_struct_fields_opt(&self) -> core::option::Option<&vortex_array::dtype::StructFields>
 
+pub fn vortex_array::dtype::DType::as_union_variants_opt(&self) -> core::option::Option<&vortex_array::dtype::UnionVariants>
+
 pub fn vortex_array::dtype::DType::element_size(&self) -> core::option::Option<usize>
 
 pub fn vortex_array::dtype::DType::eq_ignore_nullability(&self, &Self) -> bool
@@ -9396,6 +9400,8 @@ pub fn vortex_array::dtype::DType::into_struct_fields(self) -> vortex_array::dty
 
 pub fn vortex_array::dtype::DType::into_struct_fields_opt(self) -> core::option::Option<vortex_array::dtype::StructFields>
 
+pub fn vortex_array::dtype::DType::into_union_variants_opt(self) -> core::option::Option<vortex_array::dtype::UnionVariants>
+
 pub fn vortex_array::dtype::DType::is_binary(&self) -> bool
 
 pub fn vortex_array::dtype::DType::is_boolean(&self) -> bool
@@ -9422,6 +9428,8 @@ pub fn vortex_array::dtype::DType::is_signed_int(&self) -> bool
 
 pub fn vortex_array::dtype::DType::is_struct(&self) -> bool
 
+pub fn vortex_array::dtype::DType::is_union(&self) -> bool
+
 pub fn vortex_array::dtype::DType::is_unsigned_int(&self) -> bool
 
 pub fn vortex_array::dtype::DType::is_utf8(&self) -> bool
@@ -10506,6 +10514,66 @@ impl<T, V> core::iter::traits::collect::FromIterator<(T, V)> for vortex_array::d
 
 pub fn vortex_array::dtype::StructFields::from_iter<I: core::iter::traits::collect::IntoIterator<Item = (T, V)>>(I) -> Self
 
+pub struct vortex_array::dtype::UnionVariants(_)
+
+impl vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::child_index_to_tag(&self, usize) -> i8
+
+pub fn vortex_array::dtype::UnionVariants::empty() -> Self
+
+pub fn vortex_array::dtype::UnionVariants::find(&self, impl core::convert::AsRef<str>) -> core::option::Option<usize>
+
+pub fn vortex_array::dtype::UnionVariants::is_consecutive(&self) -> bool
+
+pub fn vortex_array::dtype::UnionVariants::is_empty(&self) -> bool
+
+pub fn vortex_array::dtype::UnionVariants::len(&self) -> usize
+
+pub fn vortex_array::dtype::UnionVariants::names(&self) -> &vortex_array::dtype::FieldNames
+
+pub fn vortex_array::dtype::UnionVariants::new_consecutive(vortex_array::dtype::FieldNames, alloc::vec::Vec<vortex_array::dtype::DType>) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_array::dtype::UnionVariants::nullability_constraints_satisfied(&self, vortex_array::dtype::Nullability) -> bool
+
+pub fn vortex_array::dtype::UnionVariants::tag_to_child_index(&self, i8) -> core::option::Option<usize>
+
+pub fn vortex_array::dtype::UnionVariants::try_new(vortex_array::dtype::FieldNames, alloc::vec::Vec<vortex_array::dtype::DType>, alloc::vec::Vec<i8>) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_array::dtype::UnionVariants::type_ids(&self) -> &[i8]
+
+pub fn vortex_array::dtype::UnionVariants::variant(&self, impl core::convert::AsRef<str>) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::dtype::UnionVariants::variant_by_index(&self, usize) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::dtype::UnionVariants::variants(&self) -> impl core::iter::traits::exact_size::ExactSizeIterator<Item = vortex_array::dtype::DType> + '_
+
+impl core::clone::Clone for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::clone(&self) -> vortex_array::dtype::UnionVariants
+
+impl core::cmp::Eq for vortex_array::dtype::UnionVariants
+
+impl core::cmp::PartialEq for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::eq(&self, &Self) -> bool
+
+impl core::default::Default for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::default() -> Self
+
+impl core::fmt::Debug for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::fmt::Display for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::hash::Hash for vortex_array::dtype::UnionVariants
+
+pub fn vortex_array::dtype::UnionVariants::hash<__H: core::hash::Hasher>(&self, &mut __H)
+
 #[repr(transparent)] pub struct vortex_array::dtype::i256(_)
 
 impl vortex_array::dtype::i256

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

@@ -247,6 +247,7 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
             let canonical = array.array().clone().execute::<Canonical>(ctx)?;
             return canonical_uncompressed_size_in_bytes(&canonical, ctx);
         }
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
     };
 
     value_size
@@ -293,6 +294,9 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
         DType::Struct(fields, _) => fields
             .fields()
             .all(|field| supports_uncompressed_size_in_bytes(&field)),
+        DType::Union(variants, _) => variants
+            .variants()
+            .all(|variant| supports_uncompressed_size_in_bytes(&variant)),
         DType::Extension(ext_dtype) => {
             supports_uncompressed_size_in_bytes(ext_dtype.storage_dtype())
         }

vortex-array/src/arrays/arbitrary.rs

@@ -185,6 +185,7 @@ fn random_array_chunk(
             .vortex_expect("operation should succeed in arbitrary impl")
             .into_array())
         }
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::List(elem_dtype, null) => random_list(u, elem_dtype, *null, chunk_len),
         DType::FixedSizeList(elem_dtype, list_size, null) => {
             random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)

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

@@ -151,6 +151,7 @@ pub(crate) fn constant_canonicalize(
                 StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
             })
         }
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())),
         DType::FixedSizeList(element_dtype, list_size, _) => {
             let value = scalar.as_list();

vortex-array/src/builders/mod.rs

@@ -274,6 +274,7 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBui
             *n,
             capacity,
         )),
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::List(dtype, n) => Box::new(ListViewBuilder::<u64, u64>::with_capacity(
             Arc::clone(dtype),
             *n,

vortex-array/src/builders/tests.rs

@@ -615,6 +615,7 @@ fn create_test_scalars_for_dtype(dtype: &DType, count: usize) -> Vec<Scalar> {
                     .collect();
                 Scalar::struct_(DType::Struct(fields.clone(), *n), field_values)
             }
+            DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
             DType::List(element_dtype, n) => {
                 // Create list scalars with a few elements.
                 let elements: Vec<Scalar> = (0..=i)

vortex-array/src/canonical.rs

@@ -205,6 +205,7 @@ impl Canonical {
                     Validity::from(n),
                 )
             }),
+            DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
             DType::List(dtype, n) => Canonical::List(unsafe {
                 ListViewArray::new_unchecked(
                     Canonical::empty(dtype).into_array(),

vortex-array/src/compute/conformance/consistency.rs

@@ -705,7 +705,11 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) {
 
     // Skip non-comparable types.
     match array.dtype() {
-        DType::Null | DType::Extension(_) | DType::Struct(..) | DType::List(..) => return,
+        DType::Null
+        | DType::Extension(_)
+        | DType::Struct(..)
+        | DType::Union(..)
+        | DType::List(..) => return,
         _ => {}
     }
 
@@ -833,7 +837,11 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) {
 
     // Skip non-comparable types.
     match array.dtype() {
-        DType::Null | DType::Extension(_) | DType::Struct(..) | DType::List(..) => return,
+        DType::Null
+        | DType::Extension(_)
+        | DType::Struct(..)
+        | DType::Union(..)
+        | DType::List(..) => return,
         _ => {}
     }
 
@@ -1226,6 +1234,7 @@ fn test_cast_slice_consistency(array: &ArrayRef) {
             };
             vec![DType::Struct(fields.clone(), opposite)]
         }
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         DType::List(element_type, nullability) => {
             let opposite = match nullability {
                 Nullability::NonNullable => Nullability::Nullable,

vortex-array/src/dtype/arrow.rs

@@ -320,6 +320,7 @@ impl DType {
 
                 DataType::Struct(Fields::from(fields))
             }
+            DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
             DType::Variant(_) => vortex_bail!(
                 "DType::Variant requires Arrow Field metadata; use to_arrow_schema or a Field helper"
             ),

vortex-array/src/dtype/dtype_impl.rs

@@ -15,6 +15,7 @@ use crate::dtype::FieldDType;
 use crate::dtype::FieldName;
 use crate::dtype::PType;
 use crate::dtype::StructFields;
+use crate::dtype::UnionVariants;
 use crate::dtype::decimal::DecimalDType;
 use crate::dtype::decimal::DecimalType;
 use crate::dtype::extension::ExtDTypeRef;
@@ -62,6 +63,7 @@ impl DType {
             | Utf8(null)
             | Binary(null)
             | Struct(_, null)
+            | Union(_, null)
             | List(_, null)
             | FixedSizeList(_, _, null)
             | Variant(null) => matches!(null, Nullability::Nullable),
@@ -88,6 +90,7 @@ impl DType {
             Utf8(_) => Utf8(nullability),
             Binary(_) => Binary(nullability),
             Struct(sf, _) => Struct(sf.clone(), nullability),
+            Union(vs, _) => Union(vs.clone(), nullability),
             List(edt, _) => List(Arc::clone(edt), nullability),
             FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability),
             Extension(ext) => Extension(ext.with_nullability(nullability)),
@@ -121,6 +124,15 @@ impl DType {
                         .zip_eq(rhs_dtype.fields())
                         .all(|(l, r)| l.eq_ignore_nullability(&r)))
             }
+            (Union(lhs, _), Union(rhs, _)) => {
+                // Equal `names` implies equal length by FieldNames equality.
+                lhs.names() == rhs.names()
+                    && lhs.type_ids() == rhs.type_ids()
+                    && lhs
+                        .variants()
+                        .zip_eq(rhs.variants())
+                        .all(|(l, r)| l.eq_ignore_nullability(&r))
+            }
             (Extension(lhs_extdtype), Extension(rhs_extdtype)) => {
                 lhs_extdtype.eq_ignore_nullability(rhs_extdtype)
             }
@@ -244,6 +256,11 @@ impl DType {
         matches!(self, Struct(_, _))
     }
 
+    /// Check if `self` is a [`DType::Union`] type.
+    pub fn is_union(&self) -> bool {
+        matches!(self, Union(..))
+    }
+
     /// Check if `self` is a [`DType::Extension`] type
     pub fn is_extension(&self) -> bool {
         matches!(self, Extension(_))
@@ -258,7 +275,7 @@ impl DType {
     /// recursive type.
     pub fn is_nested(&self) -> bool {
         match self {
-            List(..) | FixedSizeList(..) | Struct(..) | Variant(..) => true,
+            List(..) | FixedSizeList(..) | Struct(..) | Union(..) | Variant(..) => true,
             Extension(ext) => ext.storage_dtype().is_nested(),
             _ => false,
         }
@@ -291,6 +308,7 @@ impl DType {
                 }
                 Some(sum)
             }
+            Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
             Extension(ext) => ext.storage_dtype().element_size(),
             Variant(_) => None,
         }
@@ -416,6 +434,24 @@ impl DType {
         }
     }
 
+    /// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise `None`.
+    pub fn as_union_variants_opt(&self) -> Option<&UnionVariants> {
+        if let Union(uv, _) = self {
+            Some(uv)
+        } else {
+            None
+        }
+    }
+
+    /// Owned version of [Self::as_union_variants_opt].
+    pub fn into_union_variants_opt(self) -> Option<UnionVariants> {
+        if let Union(uv, _) = self {
+            Some(uv)
+        } else {
+            None
+        }
+    }
+
     /// Downcast a `DType` to an `ExtDType`
     pub fn as_extension(&self) -> &ExtDTypeRef {
         let Extension(ext) = self else {
@@ -465,6 +501,7 @@ impl Display for DType {
                     .map(|(field_null, dt)| format!("{field_null}={dt}"))
                     .join(", "),
             ),
+            Union(uv, null) => write!(f, "union({uv}){null}"),
             List(edt, null) => write!(f, "list({edt}){null}"),
             FixedSizeList(edt, size, null) => write!(f, "fixed_size_list({edt})[{size}]{null}"),
             Extension(ext) => write!(f, "{}", ext),