diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index f3a49b25637..75adfa94847 100644 --- a/fuzz/src/array/fill_null.rs +++ b/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"), }) } diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index f3cecec2b10..fd8eab2d6a8 100644 --- a/fuzz/src/array/mask.rs +++ b/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"), }) } diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index e13ad75343d..887974c7d42 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/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"), }) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 490ede7f640..708c83debc4 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/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") } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 51601ce76fd..de950e59d7d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/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(_) diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 1d04a074d6f..a61a45f7dbf 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -9,6 +9,7 @@ mod list_view; mod null; mod primitive; mod struct_; +mod union; mod varbinview; use std::mem::size_of; @@ -21,6 +22,7 @@ use list_view::list_view_uncompressed_size_in_bytes; use null::null_uncompressed_size_in_bytes; use primitive::primitive_uncompressed_size_in_bytes; use struct_::struct_uncompressed_size_in_bytes; +use union::union_uncompressed_size_in_bytes; use varbinview::varbinview_uncompressed_size_in_bytes; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -199,6 +201,7 @@ 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(array) => union_uncompressed_size_in_bytes(array, ctx), Canonical::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx), Canonical::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") @@ -229,11 +232,14 @@ pub(crate) fn constant_uncompressed_size_in_bytes( array.len(), array.scalar().as_binary().value().map(|value| value.len()), )?, - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => { + DType::List(..) + | DType::FixedSizeList(..) + | DType::Struct(..) + | DType::Union(..) + | DType::Extension(_) => { let canonical = array.array().clone().execute::(ctx)?; return canonical_uncompressed_size_in_bytes(&canonical, ctx); } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") } diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs new file mode 100644 index 00000000000..2f13252de9d --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::uncompressed_size_in_bytes_u64; +use crate::ExecutionCtx; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; + +pub(super) fn union_uncompressed_size_in_bytes( + array: &UnionArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut size = uncompressed_size_in_bytes_u64(array.type_ids(), ctx)?; + + for child in array.iter_children() { + size = size + .checked_add(uncompressed_size_in_bytes_u64(child, ctx)?) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?; + } + + Ok(size) +} diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 12c3bbcd9e4..04e3f397086 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -19,11 +19,14 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::union_type_ids_dtype; use crate::arrays::variant::VariantArrayExt; use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; @@ -54,6 +57,7 @@ pub(super) fn _canonicalize( let struct_array = pack_struct_chunks(owned_chunks, ctx)?; Canonical::Struct(struct_array) } + DType::Union(..) => Canonical::Union(pack_union_chunks(owned_chunks, ctx)?), DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( &owned_chunks, array.array().validity()?, @@ -89,6 +93,46 @@ fn pack_struct_chunks(chunks: Vec, ctx: &mut ExecutionCtx) -> VortexRe .process_results(|iter| StructArray::try_concat(iter))? } +/// Packs sparse union chunks into one union with chunked type IDs and sparse children. +fn pack_union_chunks(chunks: Vec, ctx: &mut ExecutionCtx) -> VortexResult { + let union_chunks = chunks + .into_iter() + .map(|chunk| chunk.execute::(ctx)) + .collect::>>()?; + let variants = union_chunks[0].variants().clone(); + let type_ids = ChunkedArray::try_new( + union_chunks + .iter() + .map(|chunk| chunk.type_ids().clone()) + .collect(), + union_type_ids_dtype(union_chunks[0].dtype().nullability()), + )? + .into_array(); + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + ChunkedArray::try_new( + union_chunks + .iter() + .map(|chunk| { + chunk + .child(index) + .vortex_expect("variant index must have a sparse child") + .clone() + }) + .collect(), + dtype, + ) + .map(IntoArray::into_array) + }) + .collect::>>()?; + + // SAFETY: Every component was taken from structurally valid UnionArrays with the same dtype + // and chunked along identical row boundaries. + Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) }) +} + /// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children. /// /// The caller guarantees there are at least 2 chunks. diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 2eabd883742..8bb424c174f 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -251,9 +251,12 @@ impl VTable for Chunked { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { - // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path - // cannot express. - DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => { + // Nested canonical arrays need child swizzling that the builder path cannot express. + DType::Struct(..) + | DType::List(..) + | DType::FixedSizeList(..) + | DType::Union(..) + | DType::Variant(..) => { // TODO(joe)[#7674]: iterative execution here too Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) } diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..77fc9980a5d 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -8,6 +8,7 @@ use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::Canonical; use crate::ExecutionCtx; @@ -23,6 +24,7 @@ use crate::arrays::ListViewArray; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::varbinview::BinaryView; @@ -30,6 +32,8 @@ use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::DecimalType; use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::UnionVariants; use crate::match_each_decimal_value; use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; @@ -146,16 +150,22 @@ pub(crate) fn constant_canonicalize( .collect(), None => { assert!(matches!(validity, Validity::AllInvalid)); - // The struct is entirely null, so fields just need placeholder values with the - // correct dtype. We use `default_value` which returns a zero for non-nullable - // dtypes and null for nullable dtypes, preserving each field's nullability. struct_dtype .fields() .map(|dt| { - let scalar = Scalar::default_value(&dt); - ConstantArray::new(scalar, array.len()).into_array() + if array.is_empty() { + return Ok(Canonical::empty(&dt).into_array()); + } + + let scalar = try_placeholder_scalar(&dt).ok_or_else(|| { + vortex_err!( + "cannot canonicalize null constant struct: field dtype {dt} \ + has no placeholder value" + ) + })?; + Ok(ConstantArray::new(scalar, array.len()).into_array()) }) - .collect() + .collect::>>()? } }; // SAFETY: Fields are constructed from the same struct scalar, all have same @@ -164,7 +174,12 @@ pub(crate) fn constant_canonicalize( StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) }) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, nullability) => Canonical::Union(constant_canonical_union( + scalar, + variants, + *nullability, + array.len(), + )?), DType::Variant(_) => Canonical::Variant(VariantArray::try_new( array.array().clone().into_array(), None, @@ -188,6 +203,97 @@ pub(crate) fn constant_canonicalize( }) } +fn constant_canonical_union( + scalar: &Scalar, + variants: &UnionVariants, + nullability: Nullability, + len: usize, +) -> VortexResult { + if len == 0 { + return Ok(UnionArray::empty(variants.clone(), nullability)); + } + + let union = scalar.as_union(); + let selected = union.child_index().zip(union.child()); + + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + let value = match &selected { + Some((selected_child, selected_value)) if index == *selected_child => { + selected_value.clone() + } + _ => try_placeholder_scalar(&dtype).ok_or_else(|| { + vortex_err!( + "cannot canonicalize constant union: unselected variant {} ({dtype}) has \ + no placeholder value", + variants.names()[index] + ) + })?, + }; + Ok(ConstantArray::new(value, len).into_array()) + }) + .collect::>>()?; + + let type_id = match union.type_id() { + Some(type_id) => Scalar::primitive(type_id, nullability), + None => Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)), + }; + + UnionArray::try_new( + ConstantArray::new(type_id, len).into_array(), + variants.clone(), + children, + ) +} + +/// Builds an arbitrary valid scalar for array positions ignored by a parent validity or type ID. +/// This does not give the dtype a semantic default value. +fn try_placeholder_scalar(dtype: &DType) -> Option { + if let Some(default) = Scalar::try_default_value(dtype) { + return Some(default); + } + + match dtype { + DType::FixedSizeList(element_dtype, size, nullability) => { + let element = try_placeholder_scalar(element_dtype)?; + Some(Scalar::fixed_size_list( + Arc::clone(element_dtype), + vec![element; *size as usize], + *nullability, + )) + } + DType::Struct(fields, _) => { + let children = fields + .fields() + .map(|field| try_placeholder_scalar(&field)) + .collect::>>()?; + Some(Scalar::struct_(dtype.clone(), children)) + } + DType::Union(variants, nullability) => { + variants + .variants() + .enumerate() + .find_map(|(index, child_dtype)| { + let child = try_placeholder_scalar(&child_dtype)?; + Scalar::union( + variants.clone(), + variants.child_index_to_tag(index), + child, + *nullability, + ) + .ok() + }) + } + DType::Extension(ext_dtype) => { + let storage = try_placeholder_scalar(ext_dtype.storage_dtype())?; + Some(Scalar::extension_ref(ext_dtype.clone(), storage)) + } + _ => None, + } +} + fn constant_canonical_byte_view( scalar_bytes: Option<&[u8]>, dtype: &DType, diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 47a9a051666..b4a6cc05017 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -30,6 +30,7 @@ use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::dict::TakeExecute; use crate::arrays::dict::TakeReduce; +use crate::arrays::union::compute::take_union; use crate::arrays::variant::VariantArrayExt; /// Take from a canonical array using indices (codes), returning a new canonical array. @@ -53,6 +54,10 @@ 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(a) => { + let indices = codes.clone().into_array(); + Canonical::Union(take_union(a.as_view(), &indices)?) + } Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)), Canonical::Variant(a) => { let indices = codes.clone().into_array(); diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 52b0373ac49..b4653787a58 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -24,6 +24,7 @@ use crate::arrays::NullArray; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterArrayExt; +use crate::arrays::union::compute::filter_union; use crate::arrays::variant::VariantArrayExt; use crate::scalar::Scalar; use crate::validity::Validity; @@ -95,6 +96,13 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask)) } Canonical::Struct(a) => Canonical::Struct(struct_::filter_struct(&a, mask)), + Canonical::Union(a) => { + let mask = Mask::Values(Arc::clone(mask)); + Canonical::Union( + filter_union(a.as_view(), &mask) + .vortex_expect("UnionArray children must support filter"), + ) + } Canonical::Extension(a) => { let filtered_storage = a .storage_array() diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index 0de8be7321e..73bcb0a3bd7 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -17,6 +17,7 @@ use crate::arrays::ListViewArray; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::bool::BoolArrayExt; @@ -24,6 +25,7 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::struct_::StructArrayExt; +use crate::arrays::union::UnionArrayExt; use crate::arrays::variant::VariantArrayExt; use crate::builtins::ArrayBuiltins; use crate::executor::ExecutionCtx; @@ -50,6 +52,7 @@ 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(a) => Canonical::Union(mask_validity_union(a, validity)?), Canonical::Extension(a) => Canonical::Extension(mask_validity_extension(a, validity, ctx)?), Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity, ctx)?), }) @@ -144,6 +147,19 @@ fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult< Ok(unsafe { StructArray::new_unchecked(fields, struct_fields.clone(), len, new_validity) }) } +fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult { + let type_ids = array + .type_ids() + .clone() + .mask(validity.to_array(array.len()))?; + let variants = array.variants().clone(); + let children = array.children(); + + // SAFETY: Masking type IDs changes only outer validity. Values at newly-null positions are no + // longer selected, and every row-aligned component retains its length and dtype. + Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) }) +} + fn mask_validity_extension( array: ExtensionArray, validity: Validity, diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 001e20dee65..7c31a1b48b1 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/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; diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs new file mode 100644 index 00000000000..0772283ee77 --- /dev/null +++ b/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 { + 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 { + /// 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 + '_ { + 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) -> 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) -> 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> UnionArrayExt for T {} + +impl Array { + /// 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>, + ) -> 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>, + ) -> VortexResult { + 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>, + ) -> 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::(nullability).into_array(); + let children: Vec<_> = variants + .variants() + .map(|dtype| crate::Canonical::empty(&dtype).into_array()) + .collect(); + + Self::new(type_ids, variants, children) + } +} diff --git a/vortex-array/src/arrays/union/compute/filter.rs b/vortex-array/src/arrays/union/compute/filter.rs new file mode 100644 index 00000000000..5310a43b7f2 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/filter.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::filter::FilterReduce; +use crate::arrays::union::UnionArrayExt; + +pub(crate) fn filter_union(array: ArrayView<'_, Union>, mask: &Mask) -> VortexResult { + let type_ids = array.type_ids().filter(mask.clone())?; + let children = array + .iter_children() + .map(|child| child.filter(mask.clone())) + .collect::>>()?; + + // SAFETY: Filtering every row-aligned component by the same mask preserves the component + // dtypes and lengths required by the structural invariants. + Ok(unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) }) +} + +impl FilterReduce for Union { + fn filter(array: ArrayView<'_, Union>, mask: &Mask) -> VortexResult> { + Ok(Some(filter_union(array, mask)?.into_array())) + } +} diff --git a/vortex-array/src/arrays/union/compute/mask.rs b/vortex-array/src/arrays/union/compute/mask.rs new file mode 100644 index 00000000000..b0a23ca15db --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::fns::mask::MaskReduce; + +impl MaskReduce for Union { + fn mask(array: ArrayView<'_, Union>, mask: &ArrayRef) -> VortexResult> { + let type_ids = array.type_ids().clone().mask(mask.clone())?; + let children = array.children(); + + // SAFETY: Masking type IDs changes only outer validity. Values at newly-null positions are + // no longer selected, and every row-aligned component retains its length and dtype. + Ok(Some( + unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) } + .into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/union/compute/mod.rs b/vortex-array/src/arrays/union/compute/mod.rs new file mode 100644 index 00000000000..8130c88ceb4 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod filter; +pub(crate) use filter::filter_union; + +mod mask; + +pub(crate) mod rules; + +mod slice; + +mod take; +pub(crate) use take::take_union; diff --git a/vortex-array/src/arrays/union/compute/rules.rs b/vortex-array/src/arrays/union/compute/rules.rs new file mode 100644 index 00000000000..ce06e369cd4 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::arrays::Union; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&SliceReduceAdaptor(Union)), + ParentRuleSet::lift(&FilterReduceAdaptor(Union)), + ParentRuleSet::lift(&TakeReduceAdaptor(Union)), + ParentRuleSet::lift(&MaskReduceAdaptor(Union)), +]); diff --git a/vortex-array/src/arrays/union/compute/slice.rs b/vortex-array/src/arrays/union/compute/slice.rs new file mode 100644 index 00000000000..fbe2fea13d9 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/slice.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::slice::SliceReduce; +use crate::arrays::union::UnionArrayExt; + +impl SliceReduce for Union { + fn slice(array: ArrayView<'_, Union>, range: Range) -> VortexResult> { + let type_ids = array.type_ids().slice(range.clone())?; + let children = array + .iter_children() + .map(|child| child.slice(range.clone())) + .collect::>>()?; + + // SAFETY: Slicing every row-aligned component by the same range preserves the component + // dtypes and lengths required by the structural invariants. + Ok(Some( + unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) } + .into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs new file mode 100644 index 00000000000..905b457ef77 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::union::UnionArrayExt; + +pub(crate) fn take_union( + array: ArrayView<'_, Union>, + indices: &ArrayRef, +) -> VortexResult { + if indices.dtype().is_nullable() { + todo!("TODO(connor)[Union]: implement take with nullable indices for Union arrays") + } + + let type_ids = array.type_ids().take(indices.clone())?; + let children = array + .iter_children() + .map(|child| child.take(indices.clone())) + .collect::>>()?; + + // SAFETY: Taking every row-aligned component with the same non-null indices preserves the + // component dtypes and lengths required by the structural invariants and cannot introduce + // nulls. + Ok(unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) }) +} + +impl TakeReduce for Union { + fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { + Ok(Some(take_union(array, indices)?.into_array())) + } +} diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs new file mode 100644 index 00000000000..aaeab88f44d --- /dev/null +++ b/vortex-array/src/arrays/union/mod.rs @@ -0,0 +1,33 @@ +// 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; + +pub(crate) mod compute; + +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; diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests.rs new file mode 100644 index 00000000000..fff2eadc5ba --- /dev/null +++ b/vortex-array/src/arrays/union/tests.rs @@ -0,0 +1,471 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +use crate::Canonical; +use crate::CanonicalValidity; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::Chunked; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::UnionVariants; +use crate::scalar::Scalar; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::validity::Validity; + +fn variants() -> VortexResult { + UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + ) +} + +fn union_array() -> VortexResult { + UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) +} + +fn nullable_variants() -> VortexResult { + UnionVariants::try_new( + ["number", "optional"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::Nullable), + ], + vec![5, 9], + ) +} + +fn nullable_union_array() -> VortexResult { + UnionArray::try_new( + PrimitiveArray::from_option_iter([Some(5u8), None, Some(9), Some(9)]).into_array(), + nullable_variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 0, 0, 0]).into_array(), + PrimitiveArray::from_option_iter([Some(0i64), Some(0), None, Some(40)]).into_array(), + ], + ) +} + +#[test] +fn scalar_at_uses_type_id_indirection() -> VortexResult<()> { + let array = union_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.child_by_name("flag")?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + assert!(array.child_by_name_opt("missing").is_none()); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable,)? + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable,)? + ); + assert!(matches!(array.validity()?, Validity::NonNullable)); + + Ok(()) +} + +#[test] +fn validates_sparse_components() -> VortexResult<()> { + let children = vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ]; + + let mismatched_lengths = UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9]).into_array(), + variants()?, + children, + ); + assert!(mismatched_lengths.is_err()); + + let missing_child = UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![PrimitiveArray::from_iter([10i32, 0, 30]).into_array()], + ); + assert!(missing_child.is_err()); + + let nullable_child = UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::new(buffer![10i32, 0, 30], Validity::AllValid).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ); + assert!(nullable_child.is_err()); + + Ok(()) +} + +#[test] +#[should_panic(expected = "Unknown UnionArray type ID 7")] +fn invalid_type_id_panics_when_accessed() { + let array = UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 7, 5]).into_array(), + variants().vortex_expect("valid Union variants"), + vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) + .vortex_expect("structurally valid UnionArray"); + let mut ctx = array_session().create_execution_ctx(); + + let _scalar = array + .execute_scalar(1, &mut ctx) + .vortex_expect("UnionArray scalar access"); +} + +#[test] +fn outer_nulls_are_independent_from_inner_nulls() -> VortexResult<()> { + let variants = nullable_variants()?; + let array = nullable_union_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.dtype(), + &DType::Union(variants.clone(), Nullability::Nullable) + ); + assert_eq!( + array.validity()?.execute_mask(array.len(), &mut ctx)?, + Mask::from_iter([true, false, true, true]) + ); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::union(variants.clone(), 5, 10i32.into(), Nullability::Nullable,)? + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::Union(variants.clone(), Nullability::Nullable)) + ); + assert_eq!( + array.execute_scalar(2, &mut ctx)?, + Scalar::union( + variants, + 9, + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + Nullability::Nullable, + )? + ); + + Ok(()) +} + +#[test] +fn masking_adds_outer_nulls_only() -> VortexResult<()> { + let masked = union_array()? + .into_array() + .mask(BoolArray::from_iter([true, false, true]).into_array())?; + let mut ctx = array_session().create_execution_ctx(); + let masked = masked.execute::(&mut ctx)?; + + assert_eq!( + masked.dtype(), + &DType::Union(variants()?, Nullability::Nullable) + ); + assert_eq!( + masked.validity()?.execute_mask(masked.len(), &mut ctx)?, + Mask::from_iter([true, false, true]) + ); + assert_eq!( + masked.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::Union(variants()?, Nullability::Nullable)) + ); + assert_eq!( + masked.execute_scalar(2, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::Nullable,)? + ); + + Ok(()) +} + +#[test] +fn structural_operations_preserve_sparse_alignment() -> VortexResult<()> { + let array = union_array()?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let sliced = array.slice(1..3)?; + let filtered = array.filter(Mask::from_iter([true, false, true]))?; + let taken = array.take(PrimitiveArray::from_iter([2u32, 1]).into_array())?; + + assert_eq!( + sliced.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable,)? + ); + assert_eq!( + filtered.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? + ); + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? + ); + + Ok(()) +} + +#[test] +fn serde_roundtrip() -> VortexResult<()> { + let session = array_session(); + let mut execution_ctx = session.create_execution_ctx(); + for array in [union_array()?, nullable_union_array()?] { + let dtype = array.dtype().clone(); + let len = array.len(); + let array_ctx = ArrayContext::empty(); + let serialized = array.clone().into_array().serialize( + &array_ctx, + &session, + &SerializeOptions::default(), + )?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &session, + )?; + let decoded = decoded.as_::(); + + assert_eq!(decoded.variants(), array.variants()); + for index in 0..len { + assert_eq!( + decoded.array().execute_scalar(index, &mut execution_ctx)?, + array.execute_scalar(index, &mut execution_ctx)? + ); + } + } + + Ok(()) +} + +#[test] +fn constant_union_executes_to_sparse_union() -> VortexResult<()> { + let scalar = Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?; + let mut ctx = array_session().create_execution_ctx(); + let array = ConstantArray::new(scalar.clone(), 3) + .into_array() + .execute::(&mut ctx)?; + + assert_eq!(array.len(), 3); + assert_eq!(array.execute_scalar(2, &mut ctx)?, scalar); + + let variants = nullable_variants()?; + let inner_null = Scalar::union( + variants.clone(), + 9, + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + Nullability::Nullable, + )?; + let array = ConstantArray::new(inner_null.clone(), 3) + .into_array() + .execute::(&mut ctx)?; + assert_eq!( + array.validity()?.execute_mask(array.len(), &mut ctx)?, + Mask::new_true(3) + ); + assert_eq!(array.execute_scalar(1, &mut ctx)?, inner_null); + + let outer_null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + let array = ConstantArray::new(outer_null.clone(), 3) + .into_array() + .execute::(&mut ctx)?; + assert_eq!( + array.validity()?.execute_mask(array.len(), &mut ctx)?, + Mask::new_false(3) + ); + assert_eq!(array.execute_scalar(1, &mut ctx)?, outer_null); + + Ok(()) +} + +#[test] +fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> { + let nested_variants = UnionVariants::try_new( + ["value"].into(), + vec![DType::Primitive(PType::I64, Nullability::NonNullable)], + vec![3], + )?; + let nested_placeholder = Scalar::union( + nested_variants.clone(), + 3, + 0_i64.into(), + Nullability::NonNullable, + )?; + let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable); + let outer_variants = UnionVariants::try_new( + ["number", "nested"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + nested_dtype, + ], + vec![5, 9], + )?; + let mut ctx = array_session().create_execution_ctx(); + + let selected_number = Scalar::union( + outer_variants.clone(), + 5, + 42i32.into(), + Nullability::NonNullable, + )?; + let array = ConstantArray::new(selected_number, 2) + .into_array() + .execute::(&mut ctx)?; + assert_eq!( + array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?, + nested_placeholder + ); + + let outer_null = Scalar::null(DType::Union(outer_variants, Nullability::Nullable)); + let array = ConstantArray::new(outer_null, 2) + .into_array() + .execute::(&mut ctx)?; + assert_eq!( + array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?, + nested_placeholder + ); + + Ok(()) +} + +#[test] +fn constant_union_builds_deeply_nested_union_placeholders() -> VortexResult<()> { + let nested_variants = UnionVariants::try_new( + ["value"].into(), + vec![DType::Primitive(PType::I64, Nullability::NonNullable)], + vec![3], + )?; + let nested_placeholder = Scalar::union( + nested_variants.clone(), + 3, + 0_i64.into(), + Nullability::NonNullable, + )?; + let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable); + let wrapper_dtype = DType::struct_([("nested", nested_dtype)], Nullability::NonNullable); + let outer_variants = UnionVariants::try_new( + ["number", "wrapper"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + wrapper_dtype, + ], + vec![5, 9], + )?; + let selected_number = + Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?; + let mut ctx = array_session().create_execution_ctx(); + let array = ConstantArray::new(selected_number, 2) + .into_array() + .execute::(&mut ctx)?; + let wrapper = array + .child_by_name("wrapper")? + .execute_scalar(0, &mut ctx)?; + + assert_eq!( + wrapper.as_struct().field("nested"), + Some(nested_placeholder) + ); + + Ok(()) +} + +#[test] +fn empty_union_supports_variant_children() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["dynamic"].into(), + vec![DType::Variant(Nullability::NonNullable)], + vec![5], + )?; + let array = Canonical::empty(&DType::Union(variants, Nullability::NonNullable)).into_union(); + + assert_eq!(array.len(), 0); + assert_eq!( + array.child(0).map(|child| child.dtype()), + Some(&DType::Variant(Nullability::NonNullable)) + ); + + Ok(()) +} + +#[test] +fn canonical_validity_canonicalizes_union_type_ids() -> VortexResult<()> { + let array = UnionArray::try_new( + ConstantArray::new(Scalar::primitive(5_u8, Nullability::Nullable), 2).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10_i32, 20]).into_array(), + BoolArray::from_iter([false, true]).into_array(), + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + let Canonical::Union(array) = array.into_array().execute::(&mut ctx)?.0 + else { + return Err(vortex_err!( + "UnionArray must remain canonical Union storage" + )); + }; + + assert!(array.type_ids().is::()); + assert_eq!(array.dtype().nullability(), Nullability::Nullable); + + Ok(()) +} + +#[test] +fn chunked_union_packs_components() -> VortexResult<()> { + let first = union_array()?.into_array().slice(0..1)?; + let second = union_array()?.into_array().slice(1..3)?; + let dtype = first.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![first, second], dtype)?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + let canonical = chunked.execute::(&mut ctx)?; + + assert!(canonical.type_ids().is::()); + assert!(canonical.iter_children().all(|child| child.is::())); + assert_eq!( + canonical.execute_scalar(2, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs new file mode 100644 index 00000000000..5357b0d4aaa --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -0,0 +1,174 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::array::with_empty_buffers; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::array::CHILDREN_OFFSET; +use crate::arrays::union::array::TYPE_IDS_SLOT; +use crate::arrays::union::array::make_union_parts; +use crate::arrays::union::compute::rules::PARENT_RULES; +use crate::arrays::union::union_type_ids_dtype; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::serde::ArrayChildren; + +mod operations; +mod validate; +mod validity; + +/// A canonical [`Union`]-encoded array. +/// +/// This is similar to Arrow's "spase" union type. +pub type UnionArray = Array; + +/// The canonical Union encoding. +/// +/// This is similar to Arrow's "spase" union type. +#[derive(Clone, Debug)] +pub struct Union; + +impl VTable for Union { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.union"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let type_ids = slots + .get(TYPE_IDS_SLOT) + .and_then(Option::as_ref) + .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?; + let variant_arrays = slots + .get(CHILDREN_OFFSET..) + .unwrap_or_default() + .iter() + .enumerate() + .map(|(index, slot)| { + slot.as_ref() + .ok_or_else(|| vortex_err!("UnionArray is missing child {index}")) + }) + .collect::>>()?; + + validate::validate_union_components(type_ids, &variant_arrays, dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("UnionArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("UnionArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); + vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!( + children.len(), + CHILDREN_OFFSET + variants.len(), + "UnionArray expected {} children, found {}", + CHILDREN_OFFSET + variants.len(), + children.len() + ); + + let type_ids = children.get(TYPE_IDS_SLOT, &union_type_ids_dtype(*nullability), len)?; + let sparse_children = variants + .variants() + .enumerate() + .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len)) + .collect::>>()?; + + Ok(make_union_parts( + type_ids, + variants.clone(), + &sparse_children, + )) + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + if idx == TYPE_IDS_SLOT { + "type_ids".to_string() + } else { + array.variants().names()[idx - CHILDREN_OFFSET].to_string() + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + _array: ArrayView<'_, Self>, + _builder: &mut dyn ArrayBuilder, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + todo!("TODO(connor)[Union]: implement append_to_builder for Union arrays") + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} diff --git a/vortex-array/src/arrays/union/vtable/operations.rs b/vortex-array/src/arrays/union/vtable/operations.rs new file mode 100644 index 00000000000..d8adb83992d --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; + +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::arrays::Union; +use crate::arrays::union::UnionArrayExt; +use crate::scalar::Scalar; + +impl OperationsVTable for Union { + fn scalar_at( + array: ArrayView<'_, Union>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let type_id_scalar = array.type_ids().execute_scalar(index, ctx)?; + let Some(type_id) = type_id_scalar.as_primitive().typed_value::() else { + return Ok(Scalar::null(array.dtype().clone())); + }; + + let Some(child_index) = array.variants().tag_to_child_index(type_id) else { + vortex_panic!("Unknown UnionArray type ID {type_id}") + }; + let child = array + .child(child_index) + .ok_or_else(|| vortex_err!("UnionArray is missing child {child_index}"))?; + let child_scalar = child.execute_scalar(index, ctx)?; + + Scalar::union( + array.variants().clone(), + type_id, + child_scalar, + array.dtype().nullability(), + ) + } +} diff --git a/vortex-array/src/arrays/union/vtable/validate.rs b/vortex-array/src/arrays/union/vtable/validate.rs new file mode 100644 index 00000000000..06ab5c79c1e --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/validate.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::arrays::union::array::CHILDREN_OFFSET; +use crate::arrays::union::union_type_ids_dtype; +use crate::dtype::DType; + +pub(super) fn validate_union_components( + type_ids: &ArrayRef, + variant_arrays: &[&ArrayRef], + dtype: &DType, + len: usize, +) -> VortexResult<()> { + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!( + variant_arrays.len(), + variants.len(), + "UnionArray has {} slots but expected {}", + CHILDREN_OFFSET + variant_arrays.len(), + CHILDREN_OFFSET + variants.len() + ); + + let expected_union_type_ids_dtype = union_type_ids_dtype(*nullability); + vortex_ensure_eq!( + type_ids.dtype(), + &expected_union_type_ids_dtype, + "UnionArray type_ids must have dtype {expected_union_type_ids_dtype}, got {}", + type_ids.dtype() + ); + vortex_ensure_eq!( + type_ids.len(), + len, + "UnionArray type_ids length {} does not match outer length {len}", + type_ids.len() + ); + + for (index, (variant_dtype, child)) in variants + .variants() + .zip(variant_arrays.iter().copied()) + .enumerate() + { + vortex_ensure_eq!( + child.len(), + len, + "UnionArray child {index} length {} does not match outer length {len}", + child.len() + ); + vortex_ensure_eq!( + child.dtype(), + &variant_dtype, + "UnionArray child {index} has dtype {} but expected {variant_dtype}", + child.dtype() + ); + } + + Ok(()) +} diff --git a/vortex-array/src/arrays/union/vtable/validity.rs b/vortex-array/src/arrays/union/vtable/validity.rs new file mode 100644 index 00000000000..80e043534e5 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/validity.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::array::ArrayView; +use crate::array::ValidityVTable; +use crate::arrays::Union; +use crate::arrays::union::UnionArrayExt; +use crate::validity::Validity; + +impl ValidityVTable for Union { + fn validity(array: ArrayView<'_, Union>) -> VortexResult { + array.type_ids().validity() + } +} diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 71da52e1300..6fe2905ab9f 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -21,6 +21,7 @@ use crate::array::ArrayView; use crate::array::child_to_validity; use crate::arrays::Bool; use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::Extension; @@ -35,6 +36,8 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::StructArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::Variant; @@ -47,6 +50,7 @@ use crate::arrays::listview::ListViewDataParts; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::primitive::PrimitiveDataParts; use crate::arrays::struct_::StructDataParts; +use crate::arrays::union::UnionDataParts; use crate::arrays::varbinview::VarBinViewDataParts; use crate::arrays::variant::VariantArrayExt; use crate::dtype::DType; @@ -69,7 +73,7 @@ use crate::validity::Validity; /// /// Each `Canonical` variant has a corresponding [`DType`] variant, with the notable exception of /// [`Canonical::VarBinView`], which is the canonical encoding for both [`DType::Utf8`] and -/// [`DType::Binary`]. [`DType::Union`] does not yet have a public canonical array. +/// [`DType::Binary`]. /// /// # Laziness /// @@ -80,8 +84,9 @@ use crate::validity::Validity; /// /// # Arrow interoperability /// -/// All of the Vortex canonical encodings have an equivalent Arrow encoding that can be built -/// zero-copy, and the corresponding Arrow array types can also be built directly. +/// Vortex canonical encodings have equivalent Arrow encodings that can be built zero-copy, except +/// [`UnionArray`], whose independent top-level validity cannot be represented directly by an Arrow +/// union. The corresponding Arrow array types can also be built directly. /// /// The full list of canonical types and their equivalent Arrow array types are: /// @@ -128,6 +133,7 @@ pub enum Canonical { List(ListViewArray), FixedSizeList(FixedSizeListArray), Struct(StructArray), + Union(UnionArray), /// Canonical storage for extension dtypes, wrapping the canonical form of the storage dtype. Extension(ExtensionArray), /// Canonical storage for dynamic variant values, optionally with typed shredded paths. @@ -146,6 +152,7 @@ macro_rules! match_each_canonical { Canonical::List($ident) => $eval, Canonical::FixedSizeList($ident) => $eval, Canonical::Struct($ident) => $eval, + Canonical::Union($ident) => $eval, Canonical::Variant($ident) => $eval, Canonical::Extension($ident) => $eval, } @@ -228,10 +235,18 @@ impl Canonical { Validity::from(n), ) }), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), - DType::Variant(_) => { - vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant") + DType::Union(variants, nullability) => { + Canonical::Union(UnionArray::empty(variants.clone(), *nullability)) } + DType::Variant(_) => Canonical::Variant( + VariantArray::try_new( + ChunkedArray::try_new(vec![], dtype.clone()) + .vortex_expect("empty Variant core storage must be valid") + .into_array(), + None, + ) + .vortex_expect("empty VariantArray must be valid"), + ), DType::Extension(ext_dtype) => Canonical::Extension(ExtensionArray::new( ext_dtype.clone(), Canonical::empty(ext_dtype.storage_dtype()).into_array(), @@ -401,6 +416,24 @@ impl Canonical { } } + /// Return this canonical array as a sparse [`UnionArray`]. + pub fn as_union(&self) -> &UnionArray { + if let Canonical::Union(a) = self { + a + } else { + vortex_panic!("Cannot get UnionArray from {:?}", &self) + } + } + + /// Unwrap this canonical array as a sparse [`UnionArray`]. + pub fn into_union(self) -> UnionArray { + if let Canonical::Union(a) = self { + a + } else { + vortex_panic!("Cannot unwrap UnionArray from {:?}", &self) + } + } + pub fn as_extension(&self) -> &ExtensionArray { if let Canonical::Extension(a) = self { a @@ -652,6 +685,18 @@ impl Executable for CanonicalValidity { StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?) }))) } + Canonical::Union(union) => { + let UnionDataParts { + variants, + type_ids, + children, + } = union.into_data_parts(); + let type_ids = type_ids.execute::(ctx)?.0.into_array(); + + Ok(CanonicalValidity(Canonical::Union(unsafe { + UnionArray::new_unchecked(type_ids, variants, children) + }))) + } Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension( ExtensionArray::new( ext.ext_dtype().clone(), @@ -829,6 +874,27 @@ impl Executable for RecursiveCanonical { ) }))) } + Canonical::Union(union) => { + let UnionDataParts { + variants, + type_ids, + children, + } = union.into_data_parts(); + let type_ids = type_ids.execute::(ctx)?.0.into_array(); + let children = children + .iter() + .cloned() + .map(|child| { + child + .execute::(ctx) + .map(|canonical| canonical.0.into_array()) + }) + .collect::>>()?; + + Ok(RecursiveCanonical(Canonical::Union(unsafe { + UnionArray::new_unchecked(type_ids, variants, children) + }))) + } Canonical::Extension(ext) => Ok(RecursiveCanonical(Canonical::Extension( ExtensionArray::new( ext.ext_dtype().clone(), @@ -1003,6 +1069,18 @@ impl Executable for StructArray { } } +/// Execute the array to canonical form and unwrap as a [`UnionArray`]. +/// +/// This will panic if the array's dtype is not union. +impl Executable for UnionArray { + fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match array.try_downcast::() { + Ok(union_array) => Ok(union_array), + Err(array) => Ok(Canonical::execute(array, ctx)?.into_union()), + } + } +} + /// Execute the array to canonical form and unwrap as a [`VariantArray`]. /// /// This will panic if the array's dtype is not variant. @@ -1032,6 +1110,7 @@ pub enum CanonicalView<'a> { List(ArrayView<'a, ListView>), FixedSizeList(ArrayView<'a, FixedSizeList>), Struct(ArrayView<'a, Struct>), + Union(ArrayView<'a, Union>), Extension(ArrayView<'a, Extension>), Variant(ArrayView<'a, Variant>), } @@ -1047,6 +1126,7 @@ impl From> for Canonical { CanonicalView::List(a) => Canonical::List(a.into_owned()), CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.into_owned()), CanonicalView::Struct(a) => Canonical::Struct(a.into_owned()), + CanonicalView::Union(a) => Canonical::Union(a.into_owned()), CanonicalView::Extension(a) => Canonical::Extension(a.into_owned()), CanonicalView::Variant(a) => Canonical::Variant(a.into_owned()), } @@ -1065,6 +1145,7 @@ impl CanonicalView<'_> { CanonicalView::List(a) => a.array().clone(), CanonicalView::FixedSizeList(a) => a.array().clone(), CanonicalView::Struct(a) => a.array().clone(), + CanonicalView::Union(a) => a.array().clone(), CanonicalView::Extension(a) => a.array().clone(), CanonicalView::Variant(a) => a.array().clone(), } @@ -1083,6 +1164,7 @@ impl Matcher for AnyCanonical { || array.is::() || array.is::() || array.is::() + || array.is::() || array.is::() || array.is::() || array.is::() @@ -1102,6 +1184,8 @@ impl Matcher for AnyCanonical { Some(CanonicalView::Decimal(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::Struct(a)) + } else if let Some(a) = array.as_opt::() { + Some(CanonicalView::Union(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::List(a)) } else if let Some(a) = array.as_opt::() { diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 7de16913c6b..cc76ffac2bd 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -319,7 +319,7 @@ impl DType { } Some(sum) } - Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + Union(..) => None, Variant(_) => None, Extension(ext) => ext.storage_dtype().element_size(), } diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 1b4298d959a..1cbe2fd2895 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -40,7 +40,7 @@ //! # Built-in, Lazy, and Experimental Arrays //! //! Built-in arrays live in [`arrays`]. Some are canonical (`PrimitiveArray`, `StructArray`, -//! `VarBinViewArray`); others are utility or lazy arrays such as [`ChunkedArray`], +//! `UnionArray`, `VarBinViewArray`); others are utility or lazy arrays such as [`ChunkedArray`], //! [`ConstantArray`], [`FilterArray`], [`SliceArray`], and [`ScalarFnArray`]. //! Lazy arrays defer work so compute kernels can operate on encoded data or prune children //! before materialization. diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 20852779d42..e524bb666cf 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -20,6 +20,7 @@ use crate::ArrayView; use crate::CanonicalView; use crate::ColumnarView; use crate::ExecutionCtx; +use crate::IntoArray; use crate::arrays::Bool; use crate::arrays::Constant; use crate::arrays::Decimal; @@ -28,8 +29,11 @@ use crate::arrays::FixedSizeList; use crate::arrays::ListView; use crate::arrays::Null; use crate::arrays::Primitive; +use crate::arrays::UnionArray; use crate::arrays::VarBinView; use crate::arrays::struct_::compute::cast::struct_cast; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::union_type_ids_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::expression::Expression; @@ -185,6 +189,7 @@ fn cast_canonical( CanonicalView::List(a) => ::cast(a, dtype, ctx), CanonicalView::FixedSizeList(a) => ::cast(a, dtype, ctx), CanonicalView::Struct(a) => struct_cast(a, dtype, ctx), + CanonicalView::Union(a) => cast_union(a, dtype), CanonicalView::Extension(a) => ::cast(a, dtype), CanonicalView::Variant(_) => { vortex_bail!("Variant arrays don't support casting") @@ -192,6 +197,30 @@ fn cast_canonical( } } +fn cast_union( + array: ArrayView<'_, crate::arrays::Union>, + dtype: &DType, +) -> VortexResult> { + let DType::Union(target_variants, target_nullability) = dtype else { + return Ok(None); + }; + if array.variants() != target_variants { + return Ok(None); + } + + let type_ids = array + .type_ids() + .cast(union_type_ids_dtype(*target_nullability))?; + let children = array.children(); + + // SAFETY: Casting only the type IDs' nullability changes the outer union nullability. The + // variant schema, sparse child dtypes, and row alignment are unchanged. + Ok(Some( + unsafe { UnionArray::new_unchecked(type_ids, target_variants.clone(), children) } + .into_array(), + )) +} + /// Cast a constant array by dispatching to its [`CastReduce`] implementation. fn cast_constant(array: ArrayView, dtype: &DType) -> VortexResult> { ::cast(array, dtype) diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 98a769ceea9..406ab31bdc3 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -28,6 +28,7 @@ use crate::arrays::Null; use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::Struct; +use crate::arrays::Union; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; @@ -73,6 +74,7 @@ impl Default for ArraySession { this.register(ListView); this.register(FixedSizeList); this.register(Struct); + this.register(Union); this.register(Variant); this.register(Extension); diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..d5797b0b37f 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -14,6 +14,7 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Masked; use vortex_array::arrays::StructArray; +use vortex_array::arrays::UnionArray; use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::extension::ExtensionArrayExt; @@ -23,6 +24,7 @@ use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::masked::MaskedArraySlotsExt; use vortex_array::arrays::scalar_fn::AnyScalarFn; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrays::union::UnionArrayExt; use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -130,6 +132,18 @@ impl CascadingCompressor { )? .into_array()) } + Canonical::Union(union_array) => { + let type_ids = self.compress(union_array.type_ids(), exec_ctx)?; + let children = union_array + .iter_children() + .map(|child| self.compress(child, exec_ctx)) + .collect::>>()?; + + Ok( + UnionArray::try_new(type_ids, union_array.variants().clone(), children)? + .into_array(), + ) + } Canonical::List(list_view_array) => { if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { let list_array = list_from_list_view(list_view_array, exec_ctx)?; diff --git a/vortex-duckdb/src/exporter/canonical.rs b/vortex-duckdb/src/exporter/canonical.rs index 7ee47a40642..c6990811d6e 100644 --- a/vortex-duckdb/src/exporter/canonical.rs +++ b/vortex-duckdb/src/exporter/canonical.rs @@ -32,6 +32,9 @@ pub(crate) fn new_exporter( Canonical::List(array) => list_view::new_exporter(array, cache, ctx), Canonical::FixedSizeList(array) => fixed_size_list::new_exporter(array, cache, ctx), Canonical::Struct(array) => struct_::new_exporter(array, cache, ctx), + Canonical::Union(_) => { + vortex_bail!("TODO(connor)[Union]: implement DuckDB export for Union arrays") + } Canonical::Extension(ext) => extension::new_exporter(ext, ctx), Canonical::Variant(_) => { vortex_bail!("Variant arrays can't be exported to DuckDB")