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 0ca56bd1e9c..106abcd88a9 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,9 +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(_) => { - todo!("TODO(connor)[Union]: implement UncompressedSizeInBytes for Union arrays") - } + 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") @@ -236,7 +236,12 @@ pub(crate) fn constant_uncompressed_size_in_bytes( let canonical = array.array().clone().execute::(ctx)?; return canonical_uncompressed_size_in_bytes(&canonical, ctx); } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => { + todo!( + "TODO(connor)[Union]: support constant Union size accounting after constant Union \ + canonicalization defines inactive sparse-child placeholders" + ) + } DType::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") } @@ -342,6 +347,7 @@ mod tests { 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::builders::builder_with_capacity; @@ -350,6 +356,7 @@ mod tests { use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -539,6 +546,26 @@ mod tests { Ok(()) } + #[test] + fn union_sums_type_ids_and_sparse_children() -> VortexResult<()> { + let type_ids = PrimitiveArray::from_iter([5_u8, 9, 5]).into_array(); + let numbers = PrimitiveArray::from_iter([10_i32, 0, 30]).into_array(); + let flags = BoolArray::from_iter([false, true, false]).into_array(); + let expected = aggregate(&type_ids)? + aggregate(&numbers)? + aggregate(&flags)?; + let variants = UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + )?; + let array = UnionArray::try_new(type_ids, variants, vec![numbers, flags])?.into_array(); + + assert_eq!(aggregate(&array)?, expected); + Ok(()) + } + #[test] fn extension_matches_materialized_size() -> VortexResult<()> { let storage = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); 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/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 2eabd883742..8c4498373f4 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -251,6 +251,12 @@ impl VTable for Chunked { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { + DType::Union(..) => { + todo!( + "TODO(connor)[Union]: canonicalize chunked Union arrays by packing type IDs and \ + every sparse child along identical chunk boundaries" + ) + } // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path // cannot express. DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => { diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..f894a251ae3 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -164,7 +164,13 @@ pub(crate) fn constant_canonicalize( StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) }) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => { + todo!( + "TODO(connor)[Union]: canonicalize constant Union arrays in a focused follow-up \ + after defining placeholder values for every inactive sparse child, including \ + nested Struct and Union variants" + ) + } DType::Variant(_) => Canonical::Variant(VariantArray::try_new( array.array().clone().into_array(), None, diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index bfd5051d365..f02948b96a9 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -54,7 +54,10 @@ pub(crate) fn take_canonical( } Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)), Canonical::Union(_) => { - todo!("TODO(connor)[Union]: implement dictionary execution for Union arrays") + todo!( + "TODO(connor)[Union]: implement dictionary execution after Union take supports \ + nullable indices and outer null propagation" + ) } Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)), Canonical::Variant(a) => { diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 048b117933e..54dd0d66f33 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -28,17 +28,21 @@ use crate::arrays::variant::VariantArrayExt; use crate::scalar::Scalar; use crate::validity::Validity; +pub(crate) mod byte_compress; + +mod slice; +mod take; + mod bitbuffer; -mod bool; mod buffer; -pub(crate) mod byte_compress; + +mod bool; mod decimal; mod fixed_size_list; mod listview; mod primitive; -mod slice; mod struct_; -pub mod take; +mod union; mod varbinview; /// A helper function that lazily filters a [`Validity`] with selection mask values. @@ -95,9 +99,7 @@ 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(_) => { - todo!("TODO(connor)[Union]: implement filter for Union arrays") - } + Canonical::Union(a) => Canonical::Union(union::filter_union(&a, mask)), Canonical::Extension(a) => { let filtered_storage = a .storage_array() diff --git a/vortex-array/src/arrays/filter/execute/union.rs b/vortex-array/src/arrays/filter/execute/union.rs new file mode 100644 index 00000000000..c00e8fba71e --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/union.rs @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_mask::Mask; +use vortex_mask::MaskValues; + +use crate::ArrayRef; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; + +pub fn filter_union(array: &UnionArray, mask: &Arc) -> UnionArray { + let filter_mask = Mask::Values(Arc::clone(mask)); + + let type_ids = array + .type_ids() + .filter(filter_mask.clone()) + .vortex_expect("UnionArray type IDs are guaranteed to support filter"); + + let children: Vec = array + .iter_children() + .map(|child| { + child + .filter(filter_mask.clone()) + .vortex_expect("UnionArray children are guaranteed to support filter") + }) + .collect(); + + UnionArray::try_new(type_ids, array.variants().clone(), children) + .vortex_expect("filtered UnionArray children have consistent dtypes and lengths") +} diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index dee1d1b5c9d..14adcef59d6 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,9 +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(_) => { - todo!("TODO(connor)[Union]: implement masking for Union arrays") - } + 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)?), }) @@ -69,7 +69,7 @@ fn mask_validity_primitive( ) -> VortexResult { let ptype = array.ptype(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: validity has same length as values + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { PrimitiveArray::new_unchecked_from_handle( array.buffer_handle().clone(), @@ -81,7 +81,7 @@ fn mask_validity_primitive( fn mask_validity_decimal(array: DecimalArray, validity: Validity) -> VortexResult { let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { DecimalArray::new_unchecked_handle( array.buffer_handle().clone(), @@ -99,7 +99,7 @@ fn mask_validity_varbinview( ) -> VortexResult { let dtype = array.dtype().as_nullable(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { VarBinViewArray::new_handle_unchecked( array.views_handle().clone(), @@ -112,7 +112,7 @@ fn mask_validity_varbinview( fn mask_validity_listview(array: ListViewArray, validity: Validity) -> VortexResult { let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. let is_zctl = array.is_zero_copy_to_list(); Ok(unsafe { ListViewArray::new_unchecked( @@ -132,7 +132,7 @@ fn mask_validity_fixed_size_list( let len = array.len(); let list_size = array.list_size(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { FixedSizeListArray::new_unchecked(array.elements().clone(), list_size, new_validity, len) }) @@ -143,16 +143,28 @@ fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult< let new_validity = Validity::and(array.validity()?, validity)?; let fields = array.unmasked_fields(); let struct_fields = array.struct_fields(); - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. 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: We're only changing validity, not the data structure. + Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) }) +} + fn mask_validity_extension( array: ExtensionArray, validity: Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { - // For extension arrays, we need to mask the underlying storage + // For extension arrays, we need to mask the underlying storage. let storage = array.storage_array().clone().execute::(ctx)?; let masked_storage = mask_validity_canonical(storage, validity, ctx)?; let masked_storage = masked_storage.into_array(); diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs index 0772283ee77..3a7603c056f 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -20,6 +20,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::Union; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::dtype::PType; use crate::dtype::UnionVariants; /// The row-aligned array of type IDs selecting a union child. @@ -144,10 +145,7 @@ impl Array { children: impl Into>, ) -> VortexResult { vortex_ensure!( - matches!( - type_ids.dtype(), - DType::Primitive(crate::dtype::PType::U8, _) - ), + matches!(type_ids.dtype(), DType::Primitive(PType::U8, _)), "UnionArray type_ids must be u8, got {}", type_ids.dtype() ); 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..0f1cbf618b7 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -0,0 +1,24 @@ +// 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> { + UnionArray::try_new( + array.type_ids().clone().mask(mask.clone())?, + array.variants().clone(), + array.children(), + ) + .map(|a| Some(a.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..965f5c41e12 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(crate) mod rules; + +mod mask; +mod slice; 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..92a558fc9a9 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::arrays::Union; +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(&MaskReduceAdaptor(Union)), + ParentRuleSet::lift(&SliceReduceAdaptor(Union)), +]); + +// TODO(connor)[Union]: Register TakeReduce only once nullable indices can introduce outer Union +// nulls while preserving the sparse children's dtypes and row alignment. 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..5a03cbebc46 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/slice.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use itertools::Itertools; +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<'_, Self>, range: Range) -> VortexResult> { + let type_ids = array.type_ids().slice(range.clone())?; + let children: Vec = array + .iter_children() + .map(|child| child.slice(range.clone())) + .try_collect()?; + + Ok(Some( + UnionArray::try_new(type_ids, array.variants().clone(), children)?.into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs index e5f5a0f10e9..aaeab88f44d 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -20,6 +20,8 @@ pub use array::UnionArrayExt; pub use array::UnionDataParts; pub use vtable::UnionArray; +pub(crate) mod compute; + mod vtable; pub use vtable::Union; diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests.rs index 51643a353b5..be82654e17a 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests.rs @@ -17,6 +17,7 @@ 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; @@ -185,6 +186,54 @@ fn outer_nulls_are_independent_from_inner_nulls() -> VortexResult<()> { 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 slice_and_filter_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]))?; + + 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,)? + ); + + Ok(()) +} + #[test] fn serde_roundtrip() -> VortexResult<()> { let session = array_session(); diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index 23997c27707..5357b0d4aaa 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -24,6 +24,7 @@ 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; @@ -151,8 +152,8 @@ impl VTable for Union { } } - fn execute(_array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - todo!("TODO(connor)[Union]: implement execute for Union arrays") + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) } fn append_to_builder( @@ -162,4 +163,12 @@ impl VTable for Union { ) -> 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/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 488a646bba4..394adcc9924 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -186,7 +186,10 @@ fn cast_canonical( CanonicalView::FixedSizeList(a) => ::cast(a, dtype, ctx), CanonicalView::Struct(a) => struct_cast(a, dtype, ctx), CanonicalView::Union(_) => { - todo!("TODO(connor)[Union]: implement casting for Union arrays") + todo!( + "TODO(connor)[Union]: implement Union casting with conformance coverage for outer \ + nullability changes, including validation of nullable-to-nonnullable casts" + ) } CanonicalView::Extension(a) => ::cast(a, dtype), CanonicalView::Variant(_) => { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index cf6f593e79e..0447aef517a 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,8 +132,17 @@ impl CascadingCompressor { )? .into_array()) } - Canonical::Union(_) => { - todo!("TODO(connor)[Union]: implement compression for Union arrays") + 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() {