add DType::Union and UnionVariants

Adds the boilerplate for a new logical `Union` type.

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

Author: connortsui20

encodings/sparse/src/canonical.rs

@@ -159,6 +159,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
         }
         DType::Extension(_ext_dtype) => todo!(),
         DType::Variant(_) => vortex_bail!("Sparse canonicalization does not support Variant"),
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
     })
 }
 

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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/mod.rs

@@ -522,6 +522,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
         }
         // Currently, no support at all
         DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"),
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
     }
 }
 

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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             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::Extension(_) | DType::Variant(_) | DType::Union(..)) => {
             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
@@ -9422,6 +9426,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 +10512,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::FieldDType>
+
+pub fn vortex_array::dtype::UnionVariants::variant_by_index(&self, usize) -> core::option::Option<&vortex_array::dtype::FieldDType>
+
+pub fn vortex_array::dtype::UnionVariants::variants(&self) -> &[vortex_array::dtype::FieldDType]
+
+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
@@ -303,6 +304,7 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
         | DType::Decimal(..)
         | DType::Utf8(_)
         | DType::Binary(_) => true,
+        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
     }
 }