From abefbe3917b537d9eba75d3198bc769f65152f80 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 13 Jul 2026 11:24:50 -0400 Subject: [PATCH 1/6] Add Union scalar support Signed-off-by: Connor Tsui --- .../src/arrays/constant/vtable/operations.rs | 44 +++ vortex-array/src/scalar/arbitrary.rs | 17 +- vortex-array/src/scalar/cast.rs | 5 +- vortex-array/src/scalar/constructor.rs | 44 +++ vortex-array/src/scalar/display.rs | 37 ++- vortex-array/src/scalar/downcast.rs | 39 +++ vortex-array/src/scalar/mod.rs | 4 +- vortex-array/src/scalar/proto.rs | 135 ++++++++ vortex-array/src/scalar/scalar_impl.rs | 67 +++- vortex-array/src/scalar/scalar_value.rs | 91 ++++-- vortex-array/src/scalar/tests/casting.rs | 65 ++++ vortex-array/src/scalar/tests/primitives.rs | 168 +++++++++- vortex-array/src/scalar/typed_view/mod.rs | 2 + vortex-array/src/scalar/typed_view/union.rs | 309 ++++++++++++++++++ vortex-array/src/scalar/validate.rs | 76 ++++- vortex-jni/src/writer.rs | 10 +- vortex-proto/proto/scalar.proto | 7 + vortex-proto/src/generated/vortex.scalar.rs | 15 +- 18 files changed, 1092 insertions(+), 43 deletions(-) create mode 100644 vortex-array/src/scalar/typed_view/union.rs diff --git a/vortex-array/src/arrays/constant/vtable/operations.rs b/vortex-array/src/arrays/constant/vtable/operations.rs index fdd2dd67345..e3568a9c39f 100644 --- a/vortex-array/src/arrays/constant/vtable/operations.rs +++ b/vortex-array/src/arrays/constant/vtable/operations.rs @@ -18,3 +18,47 @@ impl OperationsVTable for Constant { Ok(array.scalar.clone()) } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ConstantArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::UnionVariants; + use crate::scalar::Scalar; + + #[test] + fn scalar_at_preserves_union_scalar() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + + let scalar = Scalar::union( + variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::Nullable, + )?; + let array = ConstantArray::new(scalar.clone(), 3).into_array(); + let mut ctx = crate::array_session().create_execution_ctx(); + + assert_eq!(array.execute_scalar(1, &mut ctx)?, scalar); + + let null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + let array = ConstantArray::new(null.clone(), 3).into_array(); + + assert_eq!(array.execute_scalar(1, &mut ctx)?, null); + + Ok(()) + } +} diff --git a/vortex-array/src/scalar/arbitrary.rs b/vortex-array/src/scalar/arbitrary.rs index 8424f74b175..49528f05a68 100644 --- a/vortex-array/src/scalar/arbitrary.rs +++ b/vortex-array/src/scalar/arbitrary.rs @@ -95,7 +95,22 @@ pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result { )), ) .vortex_expect("unable to construct random `Scalar`_"), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, nullability) => { + let child_index = u.choose_index(variants.len())?; + + let child_dtype = variants + .variant_by_index(child_index) + .vortex_expect("chosen union child index must be valid"); + let child = random_scalar(u, &child_dtype)?; + + Scalar::union( + variants.clone(), + variants.child_index_to_tag(child_index), + child, + *nullability, + ) + .vortex_expect("generated union scalar must be valid") + } DType::Variant(_) => todo!(), DType::Extension(..) => { unreachable!("Can't yet generate arbitrary scalars for ext dtype") diff --git a/vortex-array/src/scalar/cast.rs b/vortex-array/src/scalar/cast.rs index fd3678a74e6..d297a91d9bc 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -58,7 +58,10 @@ impl Scalar { DType::Binary(_) => self.as_binary().cast(target_dtype), DType::List(..) | DType::FixedSizeList(..) => self.as_list().cast(target_dtype), DType::Struct(..) => self.as_struct().cast(target_dtype), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => vortex_bail!( + "union scalar cast from {} to {target_dtype} is not supported", + self.dtype() + ), DType::Variant(_) => vortex_bail!("Variant scalars can't be cast to {target_dtype}"), DType::Extension(..) => self.as_extension().cast(target_dtype), } diff --git a/vortex-array/src/scalar/constructor.rs b/vortex-array/src/scalar/constructor.rs index 16c87ca27a6..e240c58a1a0 100644 --- a/vortex-array/src/scalar/constructor.rs +++ b/vortex-array/src/scalar/constructor.rs @@ -8,6 +8,9 @@ use std::sync::Arc; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::dtype::DType; @@ -15,6 +18,7 @@ use crate::dtype::DecimalDType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::UnionVariants; use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtDTypeRef; use crate::dtype::extension::ExtVTable; @@ -22,6 +26,7 @@ use crate::scalar::DecimalValue; use crate::scalar::PValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; +use crate::scalar::UnionValue; // TODO(connor): Really, we want `try_` constructors that return errors instead of just panic. impl Scalar { @@ -190,6 +195,45 @@ impl Scalar { .vortex_expect("unable to construct an extension `Scalar`") } + /// Creates a union scalar from a type ID and child scalar. + /// + /// The selected child scalar is stored in full, so an inner null child remains distinct from a + /// null at the outer union level. + /// + /// # Errors + /// + /// Returns an error if the type ID is not present in `variants` or if the child scalar's dtype + /// does not exactly match the selected variant dtype. + pub fn union( + variants: UnionVariants, + type_id: u8, + child: Scalar, + nullability: Nullability, + ) -> VortexResult { + let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| { + vortex_err!( + "union type ID {type_id} is not present in {:?}", + variants.type_ids() + ) + })?; + + let expected_dtype = variants + .variant_by_index(child_index) + .vortex_expect("type ID resolved to a valid child index"); + + vortex_ensure_eq!( + child.dtype(), + &expected_dtype, + "union type ID {type_id} selects child dtype {expected_dtype}, got {}", + child.dtype() + ); + + Self::try_new( + DType::Union(variants, nullability), + Some(ScalarValue::Union(UnionValue::new(type_id, child))), + ) + } + /// Creates a new variant scalar from a row-specific nested scalar. /// /// Use [`Scalar::null(DType::Variant(Nullability::Nullable))`][Scalar::null] for a top-level diff --git a/vortex-array/src/scalar/display.rs b/vortex-array/src/scalar/display.rs index 3e3aa56b8c3..36f47aea3e0 100644 --- a/vortex-array/src/scalar/display.rs +++ b/vortex-array/src/scalar/display.rs @@ -20,7 +20,7 @@ impl Display for Scalar { DType::Binary(_) => write!(f, "{}", self.as_binary()), DType::List(..) | DType::FixedSizeList(..) => write!(f, "{}", self.as_list()), DType::Struct(..) => write!(f, "{}", self.as_struct()), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => write!(f, "{}", self.as_union()), DType::Variant(_) => write!(f, "{}", self.as_variant()), DType::Extension(_) => write!(f, "{}", self.as_extension()), } @@ -30,6 +30,7 @@ impl Display for Scalar { #[cfg(test)] mod tests { use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; use crate::dtype::DType; use crate::dtype::FieldName; @@ -37,6 +38,7 @@ mod tests { use crate::dtype::Nullability::Nullable; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::UnionVariants; use crate::extension::datetime::Date; use crate::extension::datetime::Time; use crate::extension::datetime::TimeUnit; @@ -79,6 +81,39 @@ mod tests { ); } + #[test] + fn display_union() -> VortexResult<()> { + let variants = UnionVariants::new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullable), + DType::Utf8(NonNullable), + ], + )?; + + let scalar = Scalar::union( + variants.clone(), + 0, + Scalar::primitive(42_i32, Nullable), + Nullable, + )?; + + assert_eq!(format!("{scalar}"), "int(42i32)"); + let inner_null = Scalar::union( + variants.clone(), + 0, + Scalar::null(DType::Primitive(PType::I32, Nullable)), + Nullable, + )?; + assert_eq!(format!("{inner_null}"), "int(null)"); + assert_eq!( + format!("{}", Scalar::null(DType::Union(variants, Nullable))), + "null" + ); + + Ok(()) + } + #[test] fn display_utf8() { assert_eq!( diff --git a/vortex-array/src/scalar/downcast.rs b/vortex-array/src/scalar/downcast.rs index 1276141fc35..f1be1c2f46d 100644 --- a/vortex-array/src/scalar/downcast.rs +++ b/vortex-array/src/scalar/downcast.rs @@ -8,6 +8,7 @@ use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::vortex_panic; +use crate::dtype::DType; use crate::scalar::BinaryScalar; use crate::scalar::BoolScalar; use crate::scalar::DecimalScalar; @@ -19,6 +20,8 @@ use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::scalar::StructScalar; +use crate::scalar::UnionScalar; +use crate::scalar::UnionValue; use crate::scalar::Utf8Scalar; use crate::scalar::VariantScalar; @@ -156,6 +159,26 @@ impl Scalar { Some(ExtScalar::new_unchecked(self.dtype(), self.value())) } + /// Returns a view of the scalar as a union scalar. + /// + /// # Panics + /// + /// Panics if the scalar does not have a [`Union`](crate::dtype::DType::Union) type. + pub fn as_union(&self) -> UnionScalar<'_> { + self.as_union_opt() + .vortex_expect("Failed to convert scalar to union") + } + + /// Returns a view of the scalar as a union scalar if it has a union type. + pub fn as_union_opt(&self) -> Option> { + if !matches!(self.dtype(), DType::Union(..)) { + return None; + } + + // Scalar construction has already validated the value against this union dtype. + Some(UnionScalar::new_unchecked(self.dtype(), self.value())) + } + /// Returns a view of the scalar as a variant scalar. /// /// # Panics @@ -273,6 +296,22 @@ impl ScalarValue { } } + /// Returns the union value, panicking if the value is not a [`Union`](ScalarValue::Union). + pub fn as_union(&self) -> &UnionValue { + match self { + ScalarValue::Union(value) => value, + _ => vortex_panic!("ScalarValue is not a Union"), + } + } + + /// Returns the union value, panicking if the value is not a [`Union`](ScalarValue::Union). + pub fn into_union(self) -> UnionValue { + match self { + ScalarValue::Union(value) => value, + _ => vortex_panic!("ScalarValue is not a Union"), + } + } + /// Returns the row-specific scalar wrapped by a variant, panicking if the value is not a /// variant. pub fn as_variant(&self) -> &Scalar { diff --git a/vortex-array/src/scalar/mod.rs b/vortex-array/src/scalar/mod.rs index dbc1e17d04c..4a45e202ad5 100644 --- a/vortex-array/src/scalar/mod.rs +++ b/vortex-array/src/scalar/mod.rs @@ -4,8 +4,8 @@ //! Scalar values and types for the Vortex system. //! //! This crate provides scalar types and values that can be used to represent individual data -//! elements in the Vortex array system. [`Scalar`]s are composed of a logical data type ([`DType`]) -//! and an optional (encoding nullability) value ([`ScalarValue`]). +//! elements in the Vortex array system. A [`Scalar`] pairs a logical data type ([`DType`]) with an +//! optional non-null value ([`ScalarValue`]); [`None`] represents a null scalar. //! //! Note that the implementations of `Scalar` are split into several different modules. //! diff --git a/vortex-array/src/scalar/proto.rs b/vortex-array/src/scalar/proto.rs index f172d5491f3..909c958aca2 100644 --- a/vortex-array/src/scalar/proto.rs +++ b/vortex-array/src/scalar/proto.rs @@ -15,6 +15,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_proto::scalar as pb; use vortex_proto::scalar::ListValue; +use vortex_proto::scalar::UnionValue as PbUnionValue; use vortex_proto::scalar::scalar_value::Kind; use vortex_session::VortexSession; @@ -26,6 +27,7 @@ use crate::scalar::DecimalValue; use crate::scalar::PValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; +use crate::scalar::UnionValue; //////////////////////////////////////////////////////////////////////////////////////////////////// // Serialize INTO proto. @@ -110,6 +112,12 @@ impl From<&ScalarValue> for pb::ScalarValue { kind: Some(Kind::ListValue(ListValue { values })), } } + ScalarValue::Union(v) => pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: u32::from(v.type_id()), + value: Some(Box::new(pb::Scalar::from(v.value()))), + }))), + }, ScalarValue::Variant(v) => pb::ScalarValue { kind: Some(Kind::VariantValue(Box::new(pb::Scalar::from(v.as_ref())))), }, @@ -258,6 +266,7 @@ impl ScalarValue { Kind::StringValue(s) => Some(string_from_proto(s, dtype)?), Kind::BytesValue(b) => Some(bytes_from_proto(b, dtype)?), Kind::ListValue(v) => Some(list_from_proto(v, dtype, session)?), + Kind::UnionValue(v) => Some(union_from_proto(v, dtype, session)?), Kind::VariantValue(v) => match dtype { DType::Variant(_) => Some(ScalarValue::Variant(Box::new(Scalar::from_proto( v, session, @@ -461,6 +470,46 @@ fn list_from_proto( Ok(ScalarValue::Tuple(values)) } +/// Deserialize a present union scalar value. +fn union_from_proto( + value: &PbUnionValue, + dtype: &DType, + session: &VortexSession, +) -> VortexResult { + let DType::Union(variants, _) = dtype else { + vortex_bail!(Serde: "expected Union dtype for UnionValue, got {dtype}"); + }; + + let type_id = u8::try_from(value.type_id).map_err( + |_| vortex_err!(Serde: "union type ID {} is outside the u8 range", value.type_id), + )?; + + let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| { + vortex_err!( + Serde: "union type ID {type_id} is not present in {:?}", + variants.type_ids() + ) + })?; + + let child_dtype = variants + .variant_by_index(child_index) + .ok_or_else(|| vortex_err!(Serde: "union type ID {type_id} resolved out of bounds"))?; + + let child_proto = value + .value + .as_deref() + .ok_or_else(|| vortex_err!(Serde: "UnionValue missing child value"))?; + + let child = Scalar::from_proto(child_proto, session)?; + vortex_ensure!( + child.dtype() == &child_dtype, + Serde: "union type ID {type_id} requires child dtype {child_dtype}, got {}", + child.dtype() + ); + + Ok(ScalarValue::Union(UnionValue::new(type_id, child))) +} + #[cfg(test)] mod tests { use std::f32; @@ -477,6 +526,7 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::dtype::half::f16; use crate::scalar::DecimalValue; use crate::scalar::Scalar; @@ -663,6 +713,91 @@ mod tests { ); } + #[test] + fn test_union_scalar_roundtrip() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + + round_trip(Scalar::union( + variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::Nullable, + )?); + + let inner_null = Scalar::union( + variants.clone(), + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable, + )?; + let inner_null_proto = pb::Scalar::from(&inner_null); + + assert!(matches!( + inner_null_proto + .value + .as_deref() + .and_then(|value| value.kind.as_ref()), + Some(Kind::UnionValue(_)) + )); + round_trip(inner_null); + + let outer_null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + let outer_null_proto = pb::Scalar::from(&outer_null); + assert!(matches!( + outer_null_proto + .value + .as_deref() + .and_then(|value| value.kind.as_ref()), + Some(Kind::NullValue(_)) + )); + round_trip(outer_null); + + Ok(()) + } + + #[test] + fn test_union_proto_rejects_unknown_tag_and_wrong_child_dtype() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let dtype = DType::Union(variants, Nullability::NonNullable); + + let unknown_tag = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 7, + value: Some(Box::new(pb::Scalar::from(&Scalar::primitive( + 42_i32, + Nullability::Nullable, + )))), + }))), + }; + + assert!(ScalarValue::from_proto(&unknown_tag, &dtype, &session()).is_err()); + + let wrong_child = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 5, + value: Some(Box::new(pb::Scalar::from(&Scalar::utf8( + "wrong", + Nullability::NonNullable, + )))), + }))), + }; + + assert!(ScalarValue::from_proto(&wrong_child, &dtype, &session()).is_err()); + + Ok(()) + } + #[test] fn test_backcompat_f16_serialized_as_u64() { // Backwards compatibility test for the legacy f16 serialization format. diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index edbb1fac76e..6e8d48aa808 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -100,21 +100,28 @@ impl Scalar { /// /// See [`Scalar::zero_value`] for more details about "zero" values. /// - /// For non-nullable and nested types that may need null values in their children (as of right - /// now, that is _only_ `FixedSizeList` and `Struct`), this function will provide null default - /// children. + /// For non-nullable nested types, this function recursively creates valid default children. + /// A non-nullable union selects the first variant for which a default value exists. + /// + /// # Panics + /// + /// Panics if `dtype` has no possible value, such as an empty non-nullable union. pub fn default_value(dtype: &DType) -> Self { - let value = ScalarValue::default_value(dtype); + Self::try_default_value(dtype) + .unwrap_or_else(|| vortex_panic!("{dtype} has no default value")) + } - // SAFETY: We assume that `default_value` creates a valid `ScalarValue` for the `DType`. - unsafe { Self::new_unchecked(dtype.clone(), value) } + /// Returns a valid default scalar, or [`None`] if `dtype` is uninhabited. + pub(crate) fn try_default_value(dtype: &DType) -> Option { + let value = ScalarValue::try_default_value(dtype)?; + Self::try_new(dtype.clone(), value).ok() } /// Returns a non-null zero / identity value for the given [`DType`]. /// /// # Zero Values /// - /// Here is the list of zero values for each [`DType`] (when the [`DType`] is non-nullable): + /// Here is the list of non-null zero values for each [`DType`], regardless of its nullability: /// /// - `Null`: Does not have a "zero" value /// - `Bool`: `false` @@ -127,7 +134,12 @@ impl Scalar { /// element [`DType`] /// - `Struct`: A struct where each field has a zero value, which is determined by the field /// [`DType`] + /// - `Union`: The first variant that has a non-null zero value /// - `Extension`: The zero value of the storage [`DType`] + /// + /// # Panics + /// + /// Panics if the dtype has no non-null value, such as `Null` or an empty union. pub fn zero_value(dtype: &DType) -> Self { let value = ScalarValue::zero_value(dtype); @@ -201,7 +213,7 @@ impl Scalar { .as_struct() .fields_iter() .is_some_and(|mut fields| fields.all(|f| f.is_zero() == Some(true))), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => self.as_union().value()?.is_zero()?, DType::Variant(_) => self.as_variant().is_zero()?, DType::Extension(_) => self.as_extension().to_storage_scalar().is_zero()?, }; @@ -270,7 +282,10 @@ impl Scalar { .fields_iter() .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::()) .unwrap_or_default(), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => self + .as_union() + .value() + .map_or(0, |value| 1 + value.approx_nbytes()), DType::Variant(_) => self.as_variant().value().map_or(0, Scalar::approx_nbytes), DType::Extension(_) => self.as_extension().to_storage_scalar().approx_nbytes(), } @@ -377,6 +392,16 @@ fn partial_cmp_non_null_scalar_values( (ScalarValue::Tuple(lhs), ScalarValue::Tuple(rhs)) => { partial_cmp_tuple_values(dtype, lhs, rhs) } + (ScalarValue::Union(lhs), ScalarValue::Union(rhs)) => { + if lhs.value().is_null() && rhs.value().is_null() { + return Some(Ordering::Equal); + } + if lhs.type_id() != rhs.type_id() { + return None; + } + + lhs.value().partial_cmp(rhs.value()) + } // Variant values can have a different dtype in each row, so it doesn't make sense to // compare them. (ScalarValue::Variant(_), ScalarValue::Variant(_)) => None, @@ -443,11 +468,13 @@ mod tests { use std::sync::Arc; use rstest::rstest; + use vortex_error::VortexResult; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::UnionVariants; use crate::scalar::Scalar; fn i32_scalar(value: i32) -> Scalar { @@ -590,4 +617,26 @@ mod tests { ); assert_eq!(with_non_zero.is_zero(), Some(false)); } + + #[test] + fn union_zero_skips_uninhabited_variant() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["null", "int"].into(), + vec![ + DType::Null, + DType::Primitive(PType::I32, Nullability::NonNullable), + ], + vec![0, 4], + )?; + + let scalar = Scalar::zero_value(&DType::Union(variants, Nullability::NonNullable)); + + assert_eq!(scalar.as_union().type_id(), Some(4)); + assert_eq!( + scalar.as_union().value().and_then(|value| value.is_zero()), + Some(true) + ); + + Ok(()) + } } diff --git a/vortex-array/src/scalar/scalar_value.rs b/vortex-array/src/scalar/scalar_value.rs index ebc2101fac6..f960d8c69b5 100644 --- a/vortex-array/src/scalar/scalar_value.rs +++ b/vortex-array/src/scalar/scalar_value.rs @@ -15,6 +15,7 @@ use crate::dtype::DType; use crate::scalar::DecimalValue; use crate::scalar::PValue; use crate::scalar::Scalar; +use crate::scalar::UnionValue; /// The value stored in a [`Scalar`]. /// @@ -36,6 +37,8 @@ pub enum ScalarValue { /// /// Used as the underlying representation for list, fixed-size list, and struct scalars. Tuple(Vec>), + /// A present union value carrying its selected type ID and full child scalar. + Union(UnionValue), /// A row-specific scalar wrapped by `DType::Variant`. Variant(Box), } @@ -43,8 +46,14 @@ pub enum ScalarValue { impl ScalarValue { /// Returns the zero / identity value for the given [`DType`]. pub(super) fn zero_value(dtype: &DType) -> Self { - match dtype { - DType::Null => vortex_panic!("Null dtype has no zero value"), + Self::try_zero_value(dtype) + .unwrap_or_else(|| vortex_panic!("{dtype} has no non-null zero value")) + } + + /// Returns the non-null zero value for `dtype`, or [`None`] if no such value exists. + fn try_zero_value(dtype: &DType) -> Option { + Some(match dtype { + DType::Null => return None, DType::Bool(_) => Self::Bool(false), DType::Primitive(ptype, _) => Self::Primitive(PValue::zero(ptype)), DType::Decimal(dt, ..) => Self::Decimal(DecimalValue::zero(dt)), @@ -52,40 +61,49 @@ impl ScalarValue { DType::Binary(_) => Self::Binary(ByteBuffer::empty()), DType::List(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { - let elements = (0..*size).map(|_| Some(Self::zero_value(edt))).collect(); + let elements = (0..*size) + .map(|_| Self::try_zero_value(edt).map(Some)) + .collect::>>()?; Self::Tuple(elements) } DType::Struct(fields, _) => { let field_values = fields .fields() - .map(|f| Some(Self::zero_value(&f))) - .collect(); + .map(|f| Self::try_zero_value(&f).map(Some)) + .collect::>>()?; Self::Tuple(field_values) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, _) => { + let (index, child) = + variants.variants().enumerate().find_map(|(index, dtype)| { + let value = Self::try_zero_value(&dtype)?; + let child = Scalar::try_new(dtype, Some(value)).ok()?; + Some((index, child)) + })?; + + Self::Union(UnionValue::new(variants.child_index_to_tag(index), child)) + } DType::Variant(_) => Self::Variant(Box::new(Scalar::null(DType::Null))), DType::Extension(ext_dtype) => { // Since we have no way to define a "zero" extension value (since we have no idea // what the semantics of the extension is), a best effort attempt is to just use the // zero storage value and try to make an extension scalar from that. - Self::zero_value(ext_dtype.storage_dtype()) + Self::try_zero_value(ext_dtype.storage_dtype())? } - } + }) } - /// A similar function to [`ScalarValue::zero_value`], but for nullable [`DType`]s, this returns - /// `None` instead. + /// Returns a valid default value for `dtype`, or [`None`] if the dtype is uninhabited. /// - /// For non-nullable and nested types that may need null values in their children (as of right - /// now, that is _only_ `FixedSizeList` and `Struct`), this function will provide `None` as the - /// default child values (whereas [`ScalarValue::zero_value`] would provide `Some(_)`). - pub(super) fn default_value(dtype: &DType) -> Option { + /// The outer [`Option`] distinguishes an uninhabited dtype from a nullable dtype, whose valid + /// default is represented by the inner [`None`]. + pub(super) fn try_default_value(dtype: &DType) -> Option> { if dtype.is_nullable() { - return None; + return Some(None); } - Some(match dtype { - DType::Null => vortex_panic!("Null dtype has no zero value"), + let value = match dtype { + DType::Null => return Some(None), DType::Bool(_) => Self::Bool(false), DType::Primitive(ptype, _) => Self::Primitive(PValue::zero(ptype)), DType::Decimal(dt, ..) => Self::Decimal(DecimalValue::zero(dt)), @@ -93,22 +111,48 @@ impl ScalarValue { DType::Binary(_) => Self::Binary(ByteBuffer::empty()), DType::List(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { - let elements = (0..*size).map(|_| Self::default_value(edt)).collect(); + let elements = (0..*size) + .map(|_| Self::try_default_value(edt)) + .collect::>>()?; Self::Tuple(elements) } DType::Struct(fields, _) => { - let field_values = fields.fields().map(|f| Self::default_value(&f)).collect(); + let field_values = fields + .fields() + .map(|field| Self::try_default_value(&field)) + .collect::>>()?; Self::Tuple(field_values) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, nullability) => { + let (index, child) = + variants + .variants() + .enumerate() + .find_map(|(index, child_dtype)| { + let child_value = Self::try_default_value(&child_dtype)?; + let child = Scalar::try_new(child_dtype, child_value).ok()?; + Some((index, child)) + })?; + + return Scalar::union( + variants.clone(), + variants.child_index_to_tag(index), + child, + *nullability, + ) + .ok() + .map(Scalar::into_value); + } DType::Variant(_) => Self::Variant(Box::new(Scalar::null(DType::Null))), DType::Extension(ext_dtype) => { // Since we have no way to define a "default" extension value (since we have no idea // what the semantics of the extension is), a best effort attempt is to just use the // default storage value and try to make an extension scalar from that. - Self::default_value(ext_dtype.storage_dtype())? + Self::try_default_value(ext_dtype.storage_dtype())?? } - }) + }; + + Some(Some(value)) } } @@ -156,6 +200,9 @@ impl Display for ScalarValue { } write!(f, "]") } + ScalarValue::Union(value) => { + write!(f, "union@{}({})", value.type_id(), value.value()) + } ScalarValue::Variant(value) => write!(f, "{value}"), } } diff --git a/vortex-array/src/scalar/tests/casting.rs b/vortex-array/src/scalar/tests/casting.rs index 20eb442078a..b46646f7fe8 100644 --- a/vortex-array/src/scalar/tests/casting.rs +++ b/vortex-array/src/scalar/tests/casting.rs @@ -17,6 +17,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::UnionVariants; use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtId; use crate::dtype::extension::ExtVTable; @@ -377,4 +378,68 @@ mod tests { PValue::F16(f16_value) ); } + + #[test] + fn non_identity_union_casts_are_rejected() -> VortexResult<()> { + let source_variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + + let target_variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + ], + vec![5, 9], + )?; + let target_dtype = DType::Union(target_variants, Nullability::NonNullable); + + let value = Scalar::union( + source_variants, + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::NonNullable, + )?; + + assert!(value.cast(&target_dtype).is_err()); + + Ok(()) + } + + #[test] + fn into_nullable_changes_only_outer_union() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + let scalar = Scalar::union( + variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + + let nullable = scalar.into_nullable(); + + let expected = DType::Union(variants, Nullability::Nullable); + assert_eq!(nullable.dtype(), &expected); + assert!(nullable.dtype().is_nullable()); + assert_eq!(nullable.as_union().type_id(), Some(5)); + assert_eq!( + nullable.as_union().value().map(Scalar::dtype), + Some(&DType::Primitive(PType::I32, Nullability::NonNullable)) + ); + + Ok(()) + } } diff --git a/vortex-array/src/scalar/tests/primitives.rs b/vortex-array/src/scalar/tests/primitives.rs index 5559b505533..60f99d77601 100644 --- a/vortex-array/src/scalar/tests/primitives.rs +++ b/vortex-array/src/scalar/tests/primitives.rs @@ -11,6 +11,8 @@ mod tests { use std::sync::Arc; use vortex_buffer::ByteBuffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::DType; @@ -18,6 +20,7 @@ mod tests { use crate::dtype::NativeDecimalType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; use crate::scalar::DecimalScalar; @@ -27,6 +30,20 @@ mod tests { use crate::scalar::Scalar; use crate::scalar::ScalarValue; + fn union_variants( + int_nullability: Nullability, + utf8_nullability: Nullability, + ) -> VortexResult { + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, int_nullability), + DType::Utf8(utf8_nullability), + ], + vec![5, 9], + ) + } + fn scalar_hash(scalar: &Scalar) -> u64 { let mut hasher = DefaultHasher::new(); scalar.hash(&mut hasher); @@ -66,7 +83,43 @@ mod tests { } #[test] - fn test_scalar_nbytes() { + fn default_value_for_nullable_union_is_null() -> VortexResult<()> { + let nullable = DType::Union( + union_variants(Nullability::Nullable, Nullability::NonNullable)?, + Nullability::Nullable, + ); + + assert!(Scalar::default_value(&nullable).is_null()); + + Ok(()) + } + + #[test] + fn default_value_for_non_nullable_union_selects_first_variant() -> VortexResult<()> { + let non_nullable = DType::Union( + union_variants(Nullability::NonNullable, Nullability::NonNullable)?, + Nullability::NonNullable, + ); + + let scalar = Scalar::default_value(&non_nullable); + assert_eq!(scalar.as_union().type_id(), Some(5)); + assert_eq!(scalar.as_union().value(), Some(&Scalar::from(0_i32))); + + Ok(()) + } + + #[test] + #[should_panic(expected = "has no default value")] + fn default_value_for_empty_non_nullable_union_panics() { + Scalar::default_value(&DType::Union( + UnionVariants::new(Default::default(), vec![]) + .vortex_expect("union variants must be valid"), + Nullability::NonNullable, + )); + } + + #[test] + fn test_scalar_nbytes() -> VortexResult<()> { // Test null scalar - should be 0 bytes let null_scalar = Scalar::null(DType::Null); assert_eq!(null_scalar.approx_nbytes(), 0); @@ -142,6 +195,26 @@ mod tests { Scalar::primitive(42i32, Nullability::NonNullable), ); assert_eq!(ext_scalar.approx_nbytes(), 4); // i32 storage + + // Test union scalar: one-byte type ID plus selected child. + let variants = union_variants(Nullability::Nullable, Nullability::NonNullable)?; + let union_scalar = Scalar::union( + variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::Nullable, + )?; + + assert_eq!( + union_scalar.approx_nbytes(), + size_of::() + size_of::() + ); + assert_eq!( + Scalar::null(DType::Union(variants, Nullability::Nullable)).approx_nbytes(), + 0 + ); + + Ok(()) } #[test] @@ -438,6 +511,35 @@ mod tests { assert_eq!(scalar_hash(&nullable), scalar_hash(&non_nullable)); } + #[test] + fn test_union_scalar_equality_ignores_variant_nullability() -> VortexResult<()> { + let lhs_variants = union_variants(Nullability::Nullable, Nullability::NonNullable)?; + let rhs_variants = union_variants(Nullability::NonNullable, Nullability::Nullable)?; + + let lhs = Scalar::union( + lhs_variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::Nullable, + )?; + + let rhs = Scalar::union( + rhs_variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + + assert_eq!(lhs, rhs); + + let lhs_null = Scalar::null(DType::Union(lhs_variants, Nullability::Nullable)); + let rhs_null = Scalar::null(DType::Union(rhs_variants, Nullability::Nullable)); + + assert_eq!(lhs_null, rhs_null); + + Ok(()) + } + #[test] fn test_scalar_partial_ord_incompatible_types() { let int_scalar = Scalar::primitive(42i32, Nullability::NonNullable); @@ -477,4 +579,68 @@ mod tests { assert_eq!(scalar1, scalar2); assert_ne!(scalar1, scalar3); } + + #[test] + fn test_union_type_id_is_part_of_value_identity() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["left", "right"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I32, Nullability::NonNullable), + ], + vec![3, 8], + )?; + let left = Scalar::union( + variants.clone(), + 3, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + + let right = Scalar::union( + variants, + 8, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + + assert_ne!(left, right); + assert_eq!(left.partial_cmp(&right), None); + + Ok(()) + } + + #[test] + fn selected_null_union_children_are_equal_across_variants() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + ], + vec![3, 8], + )?; + let int_null = Scalar::union( + variants.clone(), + 3, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::NonNullable, + )?; + let string_null = Scalar::union( + variants, + 8, + Scalar::null(DType::Utf8(Nullability::Nullable)), + Nullability::NonNullable, + )?; + + assert_eq!(int_null, string_null); + assert_eq!( + int_null.partial_cmp(&string_null), + Some(std::cmp::Ordering::Equal) + ); + + assert_eq!(scalar_hash(&int_null), scalar_hash(&string_null)); + + Ok(()) + } } diff --git a/vortex-array/src/scalar/typed_view/mod.rs b/vortex-array/src/scalar/typed_view/mod.rs index 639288bd8fc..0715027a4ae 100644 --- a/vortex-array/src/scalar/typed_view/mod.rs +++ b/vortex-array/src/scalar/typed_view/mod.rs @@ -22,6 +22,7 @@ mod extension; mod list; mod primitive; mod struct_; +mod union; mod utf8; mod variant; @@ -32,5 +33,6 @@ pub use extension::*; pub use list::*; pub use primitive::*; pub use struct_::*; +pub use union::*; pub use utf8::*; pub use variant::*; diff --git a/vortex-array/src/scalar/typed_view/union.rs b/vortex-array/src/scalar/typed_view/union.rs new file mode 100644 index 00000000000..dcad04aa3a7 --- /dev/null +++ b/vortex-array/src/scalar/typed_view/union.rs @@ -0,0 +1,309 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Definitions and implementations of [`UnionScalar`] and [`UnionValue`]. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; + +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::Nullability; +use crate::dtype::UnionVariants; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; + +/// The present value stored by a union scalar. +/// +/// A null union is represented by the enclosing [`Scalar`]'s value being `None`. The nested +/// [`Scalar`] is retained in full so a present union whose selected child is null remains distinct +/// from an outer null union. +#[derive(Debug, Clone)] +pub struct UnionValue { + /// The type ID selecting a variant in the enclosing [`DType::Union`]. + type_id: u8, + /// The selected variant scalar, including its dtype and inner nullability. + /// + /// This is boxed to break the recursive layout between [`Scalar`], [`ScalarValue`], and + /// [`UnionValue`]. + value: Box, +} + +impl PartialEq for UnionValue { + fn eq(&self, other: &Self) -> bool { + if self.value.is_null() && other.value.is_null() { + return true; + } + + self.type_id == other.type_id && self.value == other.value + } +} + +impl Eq for UnionValue {} + +impl Hash for UnionValue { + fn hash(&self, state: &mut H) { + let is_null = self.value.is_null(); + is_null.hash(state); + if !is_null { + self.type_id.hash(state); + self.value.hash(state); + } + } +} + +impl UnionValue { + pub(crate) fn new(type_id: u8, value: Scalar) -> Self { + Self { + type_id, + value: Box::new(value), + } + } + + /// Returns the type ID selecting the union variant. + #[inline] + pub fn type_id(&self) -> u8 { + self.type_id + } + + /// Returns the selected variant scalar, including any inner null. + #[inline] + pub fn value(&self) -> &Scalar { + &self.value + } +} + +/// A typed view into a [`DType::Union`] scalar. +/// +/// A present union scalar carries a type ID and the full scalar of the selected variant. An outer +/// null union scalar has neither because outer nullness is represented by the enclosing [`Scalar`]. +#[derive(Debug, Clone, Copy)] +pub struct UnionScalar<'a> { + /// The data type of this scalar. + dtype: &'a DType, + /// The selected union value, or [`None`] if the union scalar is null. + value: Option<&'a UnionValue>, +} + +impl Display for UnionScalar<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let Some(name) = self.variant_name() else { + return write!(f, "null"); + }; + let value = self + .value() + .vortex_expect("non-null union scalar must have a selected value"); + + write!(f, "{name}({value})") + } +} + +impl PartialEq for UnionScalar<'_> { + fn eq(&self, other: &Self) -> bool { + self.dtype.eq_ignore_nullability(other.dtype) && self.value == other.value + } +} + +impl Eq for UnionScalar<'_> {} + +impl<'a> UnionScalar<'a> { + /// Attempts to create a union scalar view from a [`DType`] and optional [`ScalarValue`]. + /// + /// # Errors + /// + /// Returns an error if `dtype` and `value` do not form a valid union scalar. This includes a + /// non-union dtype, a null value for a non-nullable union, an unknown type ID, or a value that + /// is invalid for the selected variant dtype. + pub fn try_new(dtype: &'a DType, value: Option<&'a ScalarValue>) -> VortexResult { + Scalar::validate(dtype, value)?; + + Ok(Self::new_unchecked(dtype, value)) + } + + /// Creates a union scalar view without validating the scalar value. + /// + /// # Safety + /// + /// The caller must ensure that `dtype` and `value` form a valid union scalar. + /// + /// # Panics + /// + /// Panics if `dtype` is not a union dtype or a non-null `value` is not a union value. + pub(crate) fn new_unchecked(dtype: &'a DType, value: Option<&'a ScalarValue>) -> Self { + let DType::Union(..) = dtype else { + vortex_panic!("Expected union scalar, found {dtype}") + }; + + Self { + dtype, + value: value.map(ScalarValue::as_union), + } + } + + /// Returns the data type of this union scalar. + #[inline] + pub fn dtype(&self) -> &'a DType { + self.dtype + } + + /// Returns the variants of this union scalar. + #[inline] + pub fn variants(&self) -> &'a UnionVariants { + self.dtype + .as_union_variants_opt() + .vortex_expect("UnionScalar always has union dtype") + } + + /// Returns the outer nullability of this union scalar. + #[inline] + pub fn nullability(&self) -> Nullability { + self.dtype.nullability() + } + + /// Returns true if this union scalar is null. + #[inline] + pub fn is_null(&self) -> bool { + self.value.is_none() + } + + /// Returns the selected type ID, or `None` if this scalar is null. + #[inline] + pub fn type_id(&self) -> Option { + Some(self.value?.type_id()) + } + + /// Returns the selected variant's child index, or `None` if this scalar is null. + pub fn child_index(&self) -> Option { + self.variants().tag_to_child_index(self.type_id()?) + } + + /// Returns the selected variant's name, or `None` if this scalar is null. + pub fn variant_name(&self) -> Option<&'a FieldName> { + self.variants().names().get(self.child_index()?) + } + + /// Returns the selected variant's dtype, or `None` if this scalar is null. + pub fn variant_dtype(&self) -> Option { + self.variants().variant_by_index(self.child_index()?) + } + + /// Returns the selected child scalar, or `None` if the outer union scalar is null. + pub fn value(&self) -> Option<&'a Scalar> { + Some(self.value?.value()) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::UnionScalar; + use super::UnionValue; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::UnionVariants; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + + fn variants() -> VortexResult { + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + ) + } + + #[test] + fn non_null_union_view() -> VortexResult<()> { + let child = Scalar::primitive(42_i32, Nullability::Nullable); + let scalar = Scalar::union(variants()?, 5, child.clone(), Nullability::Nullable)?; + + let union = scalar.as_union(); + assert!(!union.is_null()); + assert_eq!(union.type_id(), Some(5)); + assert_eq!(union.child_index(), Some(0)); + assert_eq!(union.variant_name().map(AsRef::as_ref), Some("int")); + assert_eq!(union.nullability(), Nullability::Nullable); + assert_eq!(union.value(), Some(&child)); + + Ok(()) + } + + #[test] + fn inner_null_is_distinct_from_outer_null() -> VortexResult<()> { + let scalar = Scalar::union( + variants()?, + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable, + )?; + + let union = scalar.as_union(); + assert!(!union.is_null()); + assert_eq!(union.type_id(), Some(5)); + assert!(union.value().is_some_and(Scalar::is_null)); + + let outer_null = Scalar::null(DType::Union(variants()?, Nullability::Nullable)); + let outer_union = outer_null.as_union(); + assert!(outer_union.is_null()); + assert_eq!(outer_union.type_id(), None); + assert_eq!(outer_union.value(), None); + + Ok(()) + } + + #[test] + fn child_is_validated() -> VortexResult<()> { + let variants = variants()?; + let child = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); + + assert!(Scalar::union(variants.clone(), 7, child, Nullability::NonNullable).is_err()); + + let wrong_child = Scalar::null(DType::Utf8(Nullability::Nullable)); + + assert!(Scalar::union(variants, 5, wrong_child, Nullability::NonNullable).is_err()); + + Ok(()) + } + + #[test] + fn try_new_validates_type_id_and_selected_value() -> VortexResult<()> { + let dtype = DType::Union(variants()?, Nullability::Nullable); + let non_nullable_dtype = DType::Union( + UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![5], + )?, + Nullability::NonNullable, + ); + + assert!(UnionScalar::try_new(&non_nullable_dtype, None).is_err()); + + let unknown_type_id = ScalarValue::Union(UnionValue::new( + 7, + Scalar::primitive(42_i32, Nullability::Nullable), + )); + + assert!(UnionScalar::try_new(&dtype, Some(&unknown_type_id)).is_err()); + + let wrong_value = ScalarValue::Union(UnionValue::new( + 5, + Scalar::utf8("wrong", Nullability::NonNullable), + )); + + assert!(UnionScalar::try_new(&dtype, Some(&wrong_value)).is_err()); + + Ok(()) + } +} diff --git a/vortex-array/src/scalar/validate.rs b/vortex-array/src/scalar/validate.rs index 1423605b1d1..8cab5f5a158 100644 --- a/vortex-array/src/scalar/validate.rs +++ b/vortex-array/src/scalar/validate.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -118,7 +119,30 @@ impl Scalar { Self::validate(&field, field_value.as_ref())?; } } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, _) => { + let ScalarValue::Union(union_value) = value else { + vortex_bail!("union dtype expected Union value, got {value}"); + }; + + let type_id = union_value.type_id(); + let Some(child_index) = variants.tag_to_child_index(type_id) else { + vortex_bail!( + "union value has unknown type ID {type_id}; expected one of {:?}", + variants.type_ids() + ); + }; + + let child_dtype = variants + .variant_by_index(child_index) + .vortex_expect("resolved union child index must be valid"); + + vortex_ensure_eq!( + union_value.value().dtype(), + &child_dtype, + "union value for type ID {type_id} must have dtype {child_dtype}, got {}", + union_value.value().dtype() + ); + } DType::Variant(_) => { let ScalarValue::Variant(inner) = value else { vortex_bail!("variant dtype expected Variant value, got {value}"); @@ -137,3 +161,53 @@ impl Scalar { Ok(()) } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::UnionVariants; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + use crate::scalar::UnionValue; + + #[test] + fn union_rejects_unknown_tag_and_wrong_value() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + let dtype = DType::Union(variants, Nullability::NonNullable); + + assert!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Union(UnionValue::new( + 7, + Scalar::primitive(42_i32, Nullability::Nullable), + ))), + ) + .is_err() + ); + + assert!( + Scalar::try_new( + dtype, + Some(ScalarValue::Union(UnionValue::new( + 5, + Scalar::utf8("wrong", Nullability::NonNullable), + ))), + ) + .is_err() + ); + + Ok(()) + } +} diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index c1ffea2f14f..53649eef4ab 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -307,10 +307,12 @@ fn scalar_to_java<'local>( } ScalarValue::Utf8(value) => Ok(env.new_string(value.as_str())?.into()), ScalarValue::Binary(value) => Ok(env.byte_array_from_slice(value.as_slice())?.into()), - ScalarValue::Tuple(_) | ScalarValue::Variant(_) => Err(JNIError::Vortex(vortex_err!( - "cannot return nested scalar write statistic with dtype {} to Java", - scalar.dtype() - ))), + ScalarValue::Tuple(_) | ScalarValue::Union(_) | ScalarValue::Variant(_) => { + Err(JNIError::Vortex(vortex_err!( + "cannot return nested scalar write statistic with dtype {} to Java", + scalar.dtype() + ))) + } } } diff --git a/vortex-proto/proto/scalar.proto b/vortex-proto/proto/scalar.proto index 251863dc3a3..e79ae3a21f7 100644 --- a/vortex-proto/proto/scalar.proto +++ b/vortex-proto/proto/scalar.proto @@ -31,9 +31,16 @@ message ScalarValue { // Variant scalars carry a row-specific nested scalar. // See RFC 0015: https://github.com/vortex-data/rfcs/blob/develop/accepted/0015-variant-type.md Scalar variant_value = 11; + UnionValue union_value = 12; } } message ListValue { repeated ScalarValue values = 1; } + +// A present union value. Outer-null unions use ScalarValue.null_value instead. +message UnionValue { + uint32 type_id = 1; + Scalar value = 2; +} diff --git a/vortex-proto/src/generated/vortex.scalar.rs b/vortex-proto/src/generated/vortex.scalar.rs index df43794acf7..8595310cfd0 100644 --- a/vortex-proto/src/generated/vortex.scalar.rs +++ b/vortex-proto/src/generated/vortex.scalar.rs @@ -8,7 +8,10 @@ pub struct Scalar { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ScalarValue { - #[prost(oneof = "scalar_value::Kind", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11")] + #[prost( + oneof = "scalar_value::Kind", + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12" + )] pub kind: ::core::option::Option, } /// Nested message and enum types in `ScalarValue`. @@ -39,6 +42,8 @@ pub mod scalar_value { /// See RFC 0015: #[prost(message, tag = "11")] VariantValue(::prost::alloc::boxed::Box), + #[prost(message, tag = "12")] + UnionValue(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -46,3 +51,11 @@ pub struct ListValue { #[prost(message, repeated, tag = "1")] pub values: ::prost::alloc::vec::Vec, } +/// A present union value. Outer-null unions use ScalarValue.null_value instead. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UnionValue { + #[prost(uint32, tag = "1")] + pub type_id: u32, + #[prost(message, optional, boxed, tag = "2")] + pub value: ::core::option::Option<::prost::alloc::boxed::Box>, +} From 4f4ca9ce5bcba21b0f899320beb344d8a7081334 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 15 Jul 2026 16:06:24 -0400 Subject: [PATCH 2/6] Add UnionArray Signed-off-by: Connor Tsui --- fuzz/src/array/fill_null.rs | 1 + fuzz/src/array/mask.rs | 1 + fuzz/src/array/scalar_at.rs | 1 + .../src/aggregate_fn/fns/is_constant/mod.rs | 3 + .../src/aggregate_fn/fns/min_max/mod.rs | 1 + .../fns/uncompressed_size_in_bytes/mod.rs | 10 +- .../fns/uncompressed_size_in_bytes/union.rs | 25 +++ .../src/arrays/chunked/vtable/canonical.rs | 44 ++++ vortex-array/src/arrays/chunked/vtable/mod.rs | 9 +- .../src/arrays/constant/vtable/canonical.rs | 46 +++- vortex-array/src/arrays/dict/execute.rs | 12 + vortex-array/src/arrays/filter/execute/mod.rs | 9 + vortex-array/src/arrays/masked/execute.rs | 4 + vortex-array/src/arrays/mod.rs | 8 +- vortex-array/src/arrays/union/array.rs | 204 +++++++++++++++++ .../src/arrays/union/compute/filter.rs | 29 +++ vortex-array/src/arrays/union/compute/mod.rs | 7 + .../src/arrays/union/compute/rules.rs | 14 ++ .../src/arrays/union/compute/slice.rs | 30 +++ vortex-array/src/arrays/union/compute/take.rs | 34 +++ vortex-array/src/arrays/union/mod.rs | 15 ++ vortex-array/src/arrays/union/tests.rs | 201 +++++++++++++++++ vortex-array/src/arrays/union/vtable/mod.rs | 207 ++++++++++++++++++ .../src/arrays/union/vtable/operations.rs | 37 ++++ .../src/arrays/union/vtable/validity.rs | 16 ++ vortex-array/src/canonical.rs | 73 +++++- vortex-array/src/dtype/dtype_impl.rs | 2 +- vortex-array/src/lib.rs | 2 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 3 + vortex-array/src/session/mod.rs | 2 + vortex-compressor/src/compressor/cascade.rs | 14 ++ vortex-duckdb/src/exporter/canonical.rs | 3 + 32 files changed, 1053 insertions(+), 14 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs create mode 100644 vortex-array/src/arrays/union/array.rs create mode 100644 vortex-array/src/arrays/union/compute/filter.rs create mode 100644 vortex-array/src/arrays/union/compute/mod.rs create mode 100644 vortex-array/src/arrays/union/compute/rules.rs create mode 100644 vortex-array/src/arrays/union/compute/slice.rs create mode 100644 vortex-array/src/arrays/union/compute/take.rs create mode 100644 vortex-array/src/arrays/union/mod.rs create mode 100644 vortex-array/src/arrays/union/tests.rs create mode 100644 vortex-array/src/arrays/union/vtable/mod.rs create mode 100644 vortex-array/src/arrays/union/vtable/operations.rs create mode 100644 vortex-array/src/arrays/union/vtable/validity.rs diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index f3a49b25637..39b7885c248 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -49,6 +49,7 @@ pub fn fill_null_canonical_array( | Canonical::List(_) | Canonical::FixedSizeList(_) | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?, + Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"), Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"), }) } diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index f3cecec2b10..5d6e411dd22 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -149,6 +149,7 @@ pub fn mask_canonical_array( .with_nullability(masked_storage.dtype().nullability()); ExtensionArray::new(ext_dtype, masked_storage).into_array() } + Canonical::Union(_) => unreachable!("Union arrays are not fuzzed yet"), 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..7b809b5da93 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/fuzz/src/array/scalar_at.rs @@ -104,6 +104,7 @@ 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(_) => unreachable!("Union arrays are not fuzzed yet"), 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..ad89513e7cd 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(_) => { + vortex_bail!("Union arrays don't support IsConstant yet") + } 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..b66b9fe792d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -420,6 +420,7 @@ impl AggregateFnVTable for MinMax { Canonical::Struct(_) | Canonical::List(_) | Canonical::FixedSizeList(_) + | Canonical::Union(_) | Canonical::Variant(_) => { vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype()) } 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..3992a38ff82 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -19,11 +19,13 @@ 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::variant::VariantArrayExt; use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; @@ -54,6 +56,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 +92,47 @@ 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(), + DType::Primitive(PType::I8, Nullability::NonNullable), + )? + .into_array(); + let children = (0..variants.len()) + .map(|index| { + let dtype = variants + .variant_by_index(index) + .vortex_expect("variant index must have a 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 validated 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..919d4607076 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_bail; 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; @@ -164,7 +166,49 @@ 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) => { + if scalar.is_null() { + vortex_bail!( + "Canonicalizing a null union scalar is not supported until null semantics are defined" + ) + } + if variants.variants().any(|variant| variant.is_nullable()) { + vortex_bail!("Canonical UnionArray children must be non-nullable") + } + + let union = scalar.as_union(); + let type_id = union + .type_id() + .vortex_expect("non-null union scalar must have a type ID"); + let child_index = union + .child_index() + .vortex_expect("validated union scalar must select a child"); + let selected_value = union + .value() + .vortex_expect("non-null union scalar must have a value"); + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + let value = if index == child_index { + selected_value.clone() + } else { + Scalar::zero_value(&dtype) + }; + ConstantArray::new(value, array.len()).into_array() + }) + .collect::>(); + + // SAFETY: The scalar's validated type ID selects `child_index`; all sparse children + // have the declared dtype and the same length. + Canonical::Union(unsafe { + UnionArray::new_unchecked( + ConstantArray::new(type_id, array.len()).into_array(), + variants.clone(), + children, + ) + }) + } 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 47a9a051666..3dc7d0f596f 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -25,6 +25,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::VariantArray; @@ -53,6 +55,7 @@ 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) => Canonical::Union(take_union(&a, codes)?), Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)), Canonical::Variant(a) => { let indices = codes.clone().into_array(); @@ -165,6 +168,15 @@ fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { .into_owned() } +fn take_union(array: &UnionArray, codes: &PrimitiveArray) -> VortexResult { + let codes_ref = codes.clone().into_array(); + let array = array.as_view(); + Ok(::take(array, &codes_ref)? + .vortex_expect("take UnionArray should be supported") + .as_::() + .into_owned()) +} + fn take_extension( array: &ExtensionArray, codes: &PrimitiveArray, diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 52b0373ac49..0dfe7ca44f8 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -21,9 +21,11 @@ use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::Filter; use crate::arrays::NullArray; +use crate::arrays::Union; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterArrayExt; +use crate::arrays::filter::FilterReduce; use crate::arrays::variant::VariantArrayExt; use crate::scalar::Scalar; use crate::validity::Validity; @@ -95,6 +97,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) => Canonical::Union( + ::filter(a.as_view(), &Mask::Values(Arc::clone(mask))) + .vortex_expect("filter UnionArray") + .vortex_expect("UnionArray filter must be supported") + .as_::() + .into_owned(), + ), 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 c4173dd0cf6..c5fc9ac90af 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::Canonical; use crate::IntoArray; @@ -50,6 +51,9 @@ pub fn mask_validity_canonical( Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?) } Canonical::Struct(a) => Canonical::Struct(mask_validity_struct(a, validity)?), + Canonical::Union(_) => { + vortex_bail!("Masking UnionArray is not supported until null semantics are defined") + } Canonical::Extension(a) => Canonical::Extension(mask_validity_extension(a, validity, ctx)?), Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity, ctx)?), }) diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..8b8f5f28d75 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 @@ -108,6 +108,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..e4dc751cba3 --- /dev/null +++ b/vortex-array/src/arrays/union/array.rs @@ -0,0 +1,204 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter::once; +use std::sync::Arc; + +use itertools::Itertools; +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::VortexSessionExecute; +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::PType; +use crate::dtype::UnionVariants; +use crate::legacy_session; + +/// 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_slots(type_ids: &ArrayRef, children: &[ArrayRef]) -> ArraySlots { + once(Some(type_ids.clone())) + .chain(children.iter().cloned().map(Some)) + .collect() +} + +/// 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 non-nullable `i8` type IDs. + 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: i8) -> Option<&ArrayRef> { + self.child(self.variants().tag_to_child_index(type_id)?) + } +} +impl> UnionArrayExt for T {} + +impl Array { + /// Construct a canonical sparse union array. + 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. + /// + /// Until Union nullability semantics are settled, the union and every child must be + /// non-nullable. + #[allow(clippy::disallowed_methods)] + pub fn try_new( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl Into>, + ) -> VortexResult { + vortex_ensure!( + type_ids.dtype() == &DType::Primitive(PType::I8, Nullability::NonNullable), + "UnionArray type_ids must be non-nullable i8, got {}", + type_ids.dtype() + ); + + let children = children.into(); + let len = type_ids.len(); + let dtype = DType::Union(variants.clone()); + let slots = make_union_slots(&type_ids, &children); + let array = Array::try_from_parts( + ArrayParts::new(Union, dtype, len, EmptyArrayData).with_slots(slots), + )?; + + // Validate data-level tags when they are host-resident. Keep scalar access fallible so an + // invalid tag from non-host or untrusted serialized input still produces an error rather + // than unchecked indexing. + if type_ids.is_host() { + let primitive = + type_ids.execute::(&mut legacy_session().create_execution_ctx())?; + for type_id in primitive.as_slice::() { + vortex_ensure!( + variants.tag_to_child_index(*type_id).is_some(), + "UnionArray type ID {type_id} is not present in {:?}", + variants.type_ids() + ); + } + } + + Ok(array) + } + + /// Construct a canonical sparse union array without validation. + /// + /// # Safety + /// + /// The caller must ensure the type IDs are non-nullable `i8` values declared by `variants`, + /// every child has the corresponding variant dtype, and all arrays have the same length. The + /// union and its children must currently be non-nullable. + pub unsafe fn new_unchecked( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl Into>, + ) -> Self { + let children = children.into(); + let len = type_ids.len(); + let slots = make_union_slots(&type_ids, &children); + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(Union, DType::Union(variants), len, EmptyArrayData) + .with_slots(slots), + ) + } + } + + /// 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, + } + } + + /// Return the sparse child for a variant name. + pub fn child_by_name(&self, name: impl AsRef) -> VortexResult<&ArrayRef> { + let name = name.as_ref(); + let index = self + .variants() + .find(name) + .ok_or_else(|| vortex_err!("Unknown UnionArray variant {name}"))?; + Ok(self + .child(index) + .vortex_expect("variant index must have a child")) + } + + /// Create an empty array for a non-nullable union dtype. + pub(crate) fn empty(variants: UnionVariants) -> Self { + assert!( + variants.variants().all(|dtype| !dtype.is_nullable()), + "Canonical UnionArray children must be non-nullable" + ); + let type_ids = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let children = variants + .variants() + .map(|dtype| crate::Canonical::empty(&dtype).into_array()) + .collect_vec(); + // SAFETY: All components are empty and use the dtypes declared by variants. + unsafe { Self::new_unchecked(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..013c2b8c3ea --- /dev/null +++ b/vortex-array/src/arrays/union/compute/filter.rs @@ -0,0 +1,29 @@ +// 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; + +impl FilterReduce for Union { + fn filter(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 all invariants. + 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..5da8b60b65d --- /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 + +mod filter; +pub(crate) mod rules; +mod slice; +mod take; 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..4666a99b92f --- /dev/null +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -0,0 +1,14 @@ +// 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; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&SliceReduceAdaptor(Union)), + ParentRuleSet::lift(&FilterReduceAdaptor(Union)), + ParentRuleSet::lift(&TakeReduceAdaptor(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..3d912130e1b --- /dev/null +++ b/vortex-array/src/arrays/union/compute/slice.rs @@ -0,0 +1,30 @@ +// 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 all 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..a9824678922 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +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; + +impl TakeReduce for Union { + fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { + if indices.dtype().is_nullable() { + vortex_bail!("Taking UnionArray with nullable indices is not supported yet") + } + + 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 all + // invariants and cannot introduce nulls. + Ok(Some( + unsafe { UnionArray::new_unchecked(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 new file mode 100644 index 00000000000..9f54572b9ec --- /dev/null +++ b/vortex-array/src/arrays/union/mod.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod array; +pub use array::UnionArrayExt; +pub use array::UnionDataParts; + +pub(crate) mod compute; + +mod vtable; +pub use vtable::Union; +pub use vtable::UnionArray; + +#[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..f22834ed06a --- /dev/null +++ b/vortex-array/src/arrays/union/tests.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +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::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([5i8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).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.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 10i32.into())? + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 9, true.into())? + ); + 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(), + ]; + + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 7, 5]).into_array(), + variants()?, + children.clone(), + ) + .is_err() + ); + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 9]).into_array(), + variants()?, + children, + ) + .is_err() + ); + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::new(buffer![10i32, 0, 30], Validity::AllValid).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) + .is_err() + ); + + 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())? + ); + assert_eq!( + filtered.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into())? + ); + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into())? + ); + + Ok(()) +} + +#[test] +fn serde_roundtrip() -> VortexResult<()> { + let session = array_session(); + let mut execution_ctx = session.create_execution_ctx(); + let array = 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())?; + 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); + + 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())? + ); + + 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..d19c7c46563 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +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_slots; +use crate::arrays::union::compute::rules::PARENT_RULES; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::serde::ArrayChildren; + +mod operations; +mod validity; + +/// A canonical sparse [`Union`]-encoded array. +pub type UnionArray = Array; + +/// The canonical sparse Union encoding. +#[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 DType::Union(variants) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure!( + !dtype.is_nullable(), + "Nullable UnionArray is not supported yet" + ); + vortex_ensure!( + variants.variants().all(|variant| !variant.is_nullable()), + "UnionArray children must be non-nullable" + ); + vortex_ensure!( + slots.len() == CHILDREN_OFFSET + variants.len(), + "UnionArray has {} slots but expected {}", + slots.len(), + CHILDREN_OFFSET + variants.len() + ); + + let type_ids = slots[TYPE_IDS_SLOT] + .as_ref() + .vortex_expect("UnionArray type_ids slot"); + vortex_ensure!( + type_ids.dtype() == &DType::Primitive(PType::I8, Nullability::NonNullable), + "UnionArray type_ids must be non-nullable i8, got {}", + type_ids.dtype() + ); + vortex_ensure!( + type_ids.len() == len, + "UnionArray type_ids length {} does not match outer length {len}", + type_ids.len() + ); + + for (index, (slot, variant_dtype)) in slots[CHILDREN_OFFSET..] + .iter() + .zip_eq(variants.variants()) + .enumerate() + { + let child = slot + .as_ref() + .ok_or_else(|| vortex_error::vortex_err!("UnionArray missing child {index}"))?; + vortex_ensure!( + child.len() == len, + "UnionArray child {index} length {} does not match outer length {len}", + child.len() + ); + vortex_ensure!( + child.dtype() == &variant_dtype, + "UnionArray child {index} has dtype {} but expected {variant_dtype}", + child.dtype() + ); + } + + Ok(()) + } + + 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) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure!( + children.len() == CHILDREN_OFFSET + variants.len(), + "UnionArray expected {} children, found {}", + CHILDREN_OFFSET + variants.len(), + children.len() + ); + + let type_ids = children.get( + TYPE_IDS_SLOT, + &DType::Primitive(PType::I8, Nullability::NonNullable), + len, + )?; + let sparse_children = variants + .variants() + .enumerate() + .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len)) + .try_collect::<_, Vec<_>, _>()?; + let slots = make_union_slots(&type_ids, &sparse_children); + + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) + } + + 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<()> { + vortex_bail!("append_to_builder is not supported for UnionArray yet") + } + + 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..debd1b4ebac --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +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 = array + .type_ids() + .execute_scalar(index, ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| vortex_err!("UnionArray type ID at index {index} is null"))?; + let child_index = array + .variants() + .tag_to_child_index(type_id) + .ok_or_else(|| vortex_err!("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) + } +} 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..0231b7a3334 --- /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::validity::Validity; + +impl ValidityVTable for Union { + fn validity(_array: ArrayView<'_, Union>) -> VortexResult { + // TODO(connor)[Union]: define how union and child nullability interact. + Ok(Validity::NonNullable) + } +} diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 71da52e1300..f8daca7e4a9 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -35,6 +35,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 +49,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 +72,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 +83,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. +/// Except for the currently experimental union encoding, 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. /// /// The full list of canonical types and their equivalent Arrow array types are: /// @@ -128,6 +132,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 +151,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,7 +234,7 @@ impl Canonical { Validity::from(n), ) }), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants) => Canonical::Union(UnionArray::empty(variants.clone())), DType::Variant(_) => { vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant") } @@ -401,6 +407,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 +676,7 @@ impl Executable for CanonicalValidity { StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?) }))) } + union @ Canonical::Union(_) => Ok(CanonicalValidity(union)), Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension( ExtensionArray::new( ext.ext_dtype().clone(), @@ -829,6 +854,28 @@ 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() + .map(|child| { + Ok(child + .clone() + .execute::(ctx)? + .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 +1050,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 +1091,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 +1107,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 +1126,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 +1145,7 @@ impl Matcher for AnyCanonical { || array.is::() || array.is::() || array.is::() + || array.is::() || array.is::() || array.is::() || array.is::() @@ -1102,6 +1165,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 52408db66b7..a0787139581 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 9d3a59eb7a0..04be12a6d03 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..1959da25a1d 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -185,6 +185,9 @@ 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(_) => { + vortex_bail!("Union arrays don't support casting yet") + } CanonicalView::Extension(a) => ::cast(a, dtype), CanonicalView::Variant(_) => { vortex_bail!("Variant arrays don't support casting") diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 2cb8414cbfd..c844a7b8015 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -27,6 +27,7 @@ use crate::arrays::Masked; use crate::arrays::Null; use crate::arrays::Primitive; use crate::arrays::Struct; +use crate::arrays::Union; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; @@ -72,6 +73,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..0bf9067250c 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!("Union arrays can't be exported to DuckDB yet") + } Canonical::Extension(ext) => extension::new_exporter(ext, ctx), Canonical::Variant(_) => { vortex_bail!("Variant arrays can't be exported to DuckDB") From a7229da55392e4ce2f938a52ab277202c4415d9f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 15 Jul 2026 16:31:52 -0400 Subject: [PATCH 3/6] Clean up UnionArray implementation Signed-off-by: Connor Tsui --- .../src/arrays/chunked/vtable/canonical.rs | 12 +- .../src/arrays/constant/vtable/canonical.rs | 88 ++++++------ vortex-array/src/arrays/dict/execute.rs | 17 +-- vortex-array/src/arrays/filter/execute/mod.rs | 17 ++- vortex-array/src/arrays/union/array.rs | 131 ++++++++++-------- .../src/arrays/union/compute/filter.rs | 23 +-- vortex-array/src/arrays/union/compute/mod.rs | 5 + vortex-array/src/arrays/union/compute/take.rs | 36 ++--- vortex-array/src/arrays/union/mod.rs | 17 ++- vortex-array/src/arrays/union/tests.rs | 49 +++---- vortex-array/src/arrays/union/vtable/mod.rs | 41 ++---- vortex-array/src/canonical.rs | 9 +- 12 files changed, 234 insertions(+), 211 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 3992a38ff82..dcf8af39ae5 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -25,6 +25,7 @@ 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::TYPE_IDS_DTYPE; use crate::arrays::union::UnionArrayExt; use crate::arrays::variant::VariantArrayExt; use crate::builders::builder_with_capacity_in; @@ -104,14 +105,13 @@ fn pack_union_chunks(chunks: Vec, ctx: &mut ExecutionCtx) -> VortexRes .iter() .map(|chunk| chunk.type_ids().clone()) .collect(), - DType::Primitive(PType::I8, Nullability::NonNullable), + TYPE_IDS_DTYPE, )? .into_array(); - let children = (0..variants.len()) - .map(|index| { - let dtype = variants - .variant_by_index(index) - .vortex_expect("variant index must have a dtype"); + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { ChunkedArray::try_new( union_chunks .iter() diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 919d4607076..ff4c1784700 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -32,6 +32,7 @@ use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::DecimalType; use crate::dtype::Nullability; +use crate::dtype::UnionVariants; use crate::match_each_decimal_value; use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; @@ -167,47 +168,7 @@ pub(crate) fn constant_canonicalize( }) } DType::Union(variants) => { - if scalar.is_null() { - vortex_bail!( - "Canonicalizing a null union scalar is not supported until null semantics are defined" - ) - } - if variants.variants().any(|variant| variant.is_nullable()) { - vortex_bail!("Canonical UnionArray children must be non-nullable") - } - - let union = scalar.as_union(); - let type_id = union - .type_id() - .vortex_expect("non-null union scalar must have a type ID"); - let child_index = union - .child_index() - .vortex_expect("validated union scalar must select a child"); - let selected_value = union - .value() - .vortex_expect("non-null union scalar must have a value"); - let children = variants - .variants() - .enumerate() - .map(|(index, dtype)| { - let value = if index == child_index { - selected_value.clone() - } else { - Scalar::zero_value(&dtype) - }; - ConstantArray::new(value, array.len()).into_array() - }) - .collect::>(); - - // SAFETY: The scalar's validated type ID selects `child_index`; all sparse children - // have the declared dtype and the same length. - Canonical::Union(unsafe { - UnionArray::new_unchecked( - ConstantArray::new(type_id, array.len()).into_array(), - variants.clone(), - children, - ) - }) + Canonical::Union(constant_canonical_union(scalar, variants, array.len())?) } DType::Variant(_) => Canonical::Variant(VariantArray::try_new( array.array().clone().into_array(), @@ -232,6 +193,51 @@ pub(crate) fn constant_canonicalize( }) } +fn constant_canonical_union( + scalar: &Scalar, + variants: &UnionVariants, + len: usize, +) -> VortexResult { + if scalar.is_null() { + vortex_bail!( + "Canonicalizing a null union scalar is not supported until null semantics are defined" + ) + } + if variants.variants().any(|variant| variant.is_nullable()) { + vortex_bail!("Canonical UnionArray children must be non-nullable") + } + + let union = scalar.as_union(); + let type_id = union + .type_id() + .vortex_expect("non-null union scalar must have a type ID"); + let selected_child = union + .child_index() + .vortex_expect("validated union scalar must select a child"); + let selected_value = union + .value() + .vortex_expect("non-null union scalar must have a value"); + + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + let value = if index == selected_child { + selected_value.clone() + } else { + Scalar::zero_value(&dtype) + }; + ConstantArray::new(value, len).into_array() + }) + .collect::>(); + + UnionArray::try_new( + ConstantArray::new(type_id, len).into_array(), + variants.clone(), + children, + ) +} + 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 3dc7d0f596f..b4a6cc05017 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -25,13 +25,12 @@ 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::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. @@ -55,7 +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) => Canonical::Union(take_union(&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(); @@ -168,15 +170,6 @@ fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { .into_owned() } -fn take_union(array: &UnionArray, codes: &PrimitiveArray) -> VortexResult { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - Ok(::take(array, &codes_ref)? - .vortex_expect("take UnionArray should be supported") - .as_::() - .into_owned()) -} - fn take_extension( array: &ExtensionArray, codes: &PrimitiveArray, diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 0dfe7ca44f8..b4653787a58 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -21,11 +21,10 @@ use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::Filter; use crate::arrays::NullArray; -use crate::arrays::Union; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterArrayExt; -use crate::arrays::filter::FilterReduce; +use crate::arrays::union::compute::filter_union; use crate::arrays::variant::VariantArrayExt; use crate::scalar::Scalar; use crate::validity::Validity; @@ -97,13 +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) => Canonical::Union( - ::filter(a.as_view(), &Mask::Values(Arc::clone(mask))) - .vortex_expect("filter UnionArray") - .vortex_expect("UnionArray filter must be supported") - .as_::() - .into_owned(), - ), + 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/union/array.rs b/vortex-array/src/arrays/union/array.rs index e4dc751cba3..fef81d2672d 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -4,7 +4,6 @@ use std::iter::once; use std::sync::Arc; -use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -20,9 +19,8 @@ use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; use crate::arrays::PrimitiveArray; use crate::arrays::Union; +use crate::arrays::union::TYPE_IDS_DTYPE; use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; use crate::dtype::UnionVariants; use crate::legacy_session; @@ -31,10 +29,17 @@ 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_slots(type_ids: &ArrayRef, children: &[ArrayRef]) -> ArraySlots { - once(Some(type_ids.clone())) +pub(super) fn make_union_parts( + type_ids: ArrayRef, + variants: UnionVariants, + children: &[ArrayRef], +) -> ArrayParts { + let len = type_ids.len(); + let slots: ArraySlots = once(Some(type_ids)) .chain(children.iter().cloned().map(Some)) - .collect() + .collect(); + + ArrayParts::new(Union, DType::Union(variants), len, EmptyArrayData).with_slots(slots) } /// Concrete parts of a [`UnionArray`](super::UnionArray). @@ -85,11 +90,56 @@ pub trait UnionArrayExt: TypedArrayRef { fn child_by_type_id(&self, type_id: i8) -> 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 {} +/// Validate host-resident type ID values eagerly. +/// +/// Non-host arrays remain valid structurally and are checked when their values are accessed. +#[allow(clippy::disallowed_methods)] +fn validate_type_id_values(array: &Array) -> VortexResult<()> { + if !array.type_ids().is_host() { + return Ok(()); + } + + let type_ids = array + .type_ids() + .clone() + .execute::(&mut legacy_session().create_execution_ctx())?; + for type_id in type_ids.as_slice::() { + vortex_ensure!( + array.variants().tag_to_child_index(*type_id).is_some(), + "UnionArray type ID {type_id} is not present in {:?}", + array.variants().type_ids() + ); + } + + Ok(()) +} + 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, @@ -100,42 +150,26 @@ impl Array { /// Try to construct a canonical sparse union array. /// - /// Until Union nullability semantics are settled, the union and every child must be - /// non-nullable. - #[allow(clippy::disallowed_methods)] + /// Until Union nullability semantics are settled, every child must be non-nullable. + /// + /// # Errors + /// + /// Returns an error if the type IDs and sparse children do not match the variant schema or do + /// not all have the same length. pub fn try_new( type_ids: ArrayRef, variants: UnionVariants, children: impl Into>, ) -> VortexResult { vortex_ensure!( - type_ids.dtype() == &DType::Primitive(PType::I8, Nullability::NonNullable), + type_ids.dtype() == &TYPE_IDS_DTYPE, "UnionArray type_ids must be non-nullable i8, got {}", type_ids.dtype() ); let children = children.into(); - let len = type_ids.len(); - let dtype = DType::Union(variants.clone()); - let slots = make_union_slots(&type_ids, &children); - let array = Array::try_from_parts( - ArrayParts::new(Union, dtype, len, EmptyArrayData).with_slots(slots), - )?; - - // Validate data-level tags when they are host-resident. Keep scalar access fallible so an - // invalid tag from non-host or untrusted serialized input still produces an error rather - // than unchecked indexing. - if type_ids.is_host() { - let primitive = - type_ids.execute::(&mut legacy_session().create_execution_ctx())?; - for type_id in primitive.as_slice::() { - vortex_ensure!( - variants.tag_to_child_index(*type_id).is_some(), - "UnionArray type ID {type_id} is not present in {:?}", - variants.type_ids() - ); - } - } + let array = Array::try_from_parts(make_union_parts(type_ids, variants, &children))?; + validate_type_id_values(&array)?; Ok(array) } @@ -146,21 +180,14 @@ impl Array { /// /// The caller must ensure the type IDs are non-nullable `i8` values declared by `variants`, /// every child has the corresponding variant dtype, and all arrays have the same length. The - /// union and its children must currently be non-nullable. + /// children must currently be non-nullable. pub unsafe fn new_unchecked( type_ids: ArrayRef, variants: UnionVariants, children: impl Into>, ) -> Self { let children = children.into(); - let len = type_ids.len(); - let slots = make_union_slots(&type_ids, &children); - unsafe { - Array::from_parts_unchecked( - ArrayParts::new(Union, DType::Union(variants), len, EmptyArrayData) - .with_slots(slots), - ) - } + unsafe { Array::from_parts_unchecked(make_union_parts(type_ids, variants, &children)) } } /// Deconstruct this array into its type IDs, variant schema, and sparse children. @@ -175,30 +202,14 @@ impl Array { } } - /// Return the sparse child for a variant name. - pub fn child_by_name(&self, name: impl AsRef) -> VortexResult<&ArrayRef> { - let name = name.as_ref(); - let index = self - .variants() - .find(name) - .ok_or_else(|| vortex_err!("Unknown UnionArray variant {name}"))?; - Ok(self - .child(index) - .vortex_expect("variant index must have a child")) - } - /// Create an empty array for a non-nullable union dtype. pub(crate) fn empty(variants: UnionVariants) -> Self { - assert!( - variants.variants().all(|dtype| !dtype.is_nullable()), - "Canonical UnionArray children must be non-nullable" - ); let type_ids = PrimitiveArray::from_iter(Vec::::new()).into_array(); - let children = variants + let children: Vec<_> = variants .variants() .map(|dtype| crate::Canonical::empty(&dtype).into_array()) - .collect_vec(); - // SAFETY: All components are empty and use the dtypes declared by variants. - unsafe { Self::new_unchecked(type_ids, variants, children) } + .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 index 013c2b8c3ea..4b26916bdbc 100644 --- a/vortex-array/src/arrays/union/compute/filter.rs +++ b/vortex-array/src/arrays/union/compute/filter.rs @@ -12,18 +12,19 @@ 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 all invariants. + Ok(unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) }) +} + impl FilterReduce for Union { fn filter(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 all invariants. - Ok(Some( - unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) } - .into_array(), - )) + Ok(Some(filter_union(array, mask)?.into_array())) } } diff --git a/vortex-array/src/arrays/union/compute/mod.rs b/vortex-array/src/arrays/union/compute/mod.rs index 5da8b60b65d..e9208bf2185 100644 --- a/vortex-array/src/arrays/union/compute/mod.rs +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -2,6 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod filter; +pub(crate) use filter::filter_union; + pub(crate) mod rules; + mod slice; + mod take; +pub(crate) use take::take_union; diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs index a9824678922..9cca9a0794a 100644 --- a/vortex-array/src/arrays/union/compute/take.rs +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -12,23 +12,27 @@ use crate::arrays::UnionArray; use crate::arrays::dict::TakeReduce; use crate::arrays::union::UnionArrayExt; -impl TakeReduce for Union { - fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { - if indices.dtype().is_nullable() { - vortex_bail!("Taking UnionArray with nullable indices is not supported yet") - } +pub(crate) fn take_union( + array: ArrayView<'_, Union>, + indices: &ArrayRef, +) -> VortexResult { + if indices.dtype().is_nullable() { + vortex_bail!("Taking UnionArray with nullable indices is not supported yet") + } - let type_ids = array.type_ids().take(indices.clone())?; - let children = array - .iter_children() - .map(|child| child.take(indices.clone())) - .collect::>>()?; + 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 all - // invariants and cannot introduce nulls. - Ok(Some( - unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) } - .into_array(), - )) + // SAFETY: Taking every row-aligned component with the same non-null indices preserves all + // 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 index 9f54572b9ec..b8671ccafd8 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -1,15 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Canonical sparse union arrays. +//! +//! A [`UnionArray`] stores one non-nullable `i8` 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. +//! +//! Union nullability semantics are still being designed, so this encoding currently requires every +//! variant child to be non-nullable. + +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 use vtable::UnionArray; + +pub(crate) const TYPE_IDS_DTYPE: DType = DType::Primitive(PType::I8, Nullability::NonNullable); #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests.rs index f22834ed06a..e800647d56b 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests.rs @@ -55,6 +55,11 @@ 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())? @@ -75,33 +80,29 @@ fn validates_sparse_components() -> VortexResult<()> { BoolArray::from_iter([false, true, false]).into_array(), ]; - assert!( - UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 7, 5]).into_array(), - variants()?, - children.clone(), - ) - .is_err() + let unknown_type_id = UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 7, 5]).into_array(), + variants()?, + children.clone(), ); - assert!( - UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 9]).into_array(), - variants()?, - children, - ) - .is_err() + assert!(unknown_type_id.is_err()); + + let mismatched_lengths = UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 9]).into_array(), + variants()?, + children, ); - assert!( - UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 9, 5]).into_array(), - variants()?, - vec![ - PrimitiveArray::new(buffer![10i32, 0, 30], Validity::AllValid).into_array(), - BoolArray::from_iter([false, true, false]).into_array(), - ], - ) - .is_err() + assert!(mismatched_lengths.is_err()); + + let nullable_child = UnionArray::try_new( + PrimitiveArray::from_iter([5i8, 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(()) } diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index d19c7c46563..4023699c256 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -1,11 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -20,16 +19,15 @@ use crate::array::ArrayView; use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::with_empty_buffers; +use crate::arrays::union::TYPE_IDS_DTYPE; 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_slots; +use crate::arrays::union::array::make_union_parts; use crate::arrays::union::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; use crate::serde::ArrayChildren; mod operations; @@ -62,10 +60,6 @@ impl VTable for Union { let DType::Union(variants) = dtype else { vortex_bail!("Expected union dtype, found {dtype}") }; - vortex_ensure!( - !dtype.is_nullable(), - "Nullable UnionArray is not supported yet" - ); vortex_ensure!( variants.variants().all(|variant| !variant.is_nullable()), "UnionArray children must be non-nullable" @@ -79,9 +73,9 @@ impl VTable for Union { let type_ids = slots[TYPE_IDS_SLOT] .as_ref() - .vortex_expect("UnionArray type_ids slot"); + .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?; vortex_ensure!( - type_ids.dtype() == &DType::Primitive(PType::I8, Nullability::NonNullable), + type_ids.dtype() == &TYPE_IDS_DTYPE, "UnionArray type_ids must be non-nullable i8, got {}", type_ids.dtype() ); @@ -91,14 +85,10 @@ impl VTable for Union { type_ids.len() ); - for (index, (slot, variant_dtype)) in slots[CHILDREN_OFFSET..] - .iter() - .zip_eq(variants.variants()) - .enumerate() - { - let child = slot + for (index, variant_dtype) in variants.variants().enumerate() { + let child = slots[CHILDREN_OFFSET + index] .as_ref() - .ok_or_else(|| vortex_error::vortex_err!("UnionArray missing child {index}"))?; + .ok_or_else(|| vortex_err!("UnionArray is missing child {index}"))?; vortex_ensure!( child.len() == len, "UnionArray child {index} length {} does not match outer length {len}", @@ -162,19 +152,18 @@ impl VTable for Union { children.len() ); - let type_ids = children.get( - TYPE_IDS_SLOT, - &DType::Primitive(PType::I8, Nullability::NonNullable), - len, - )?; + let type_ids = children.get(TYPE_IDS_SLOT, &TYPE_IDS_DTYPE, len)?; let sparse_children = variants .variants() .enumerate() .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len)) - .try_collect::<_, Vec<_>, _>()?; - let slots = make_union_slots(&type_ids, &sparse_children); + .collect::>>()?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) + Ok(make_union_parts( + type_ids, + variants.clone(), + &sparse_children, + )) } fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index f8daca7e4a9..2e4851ac15d 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -863,12 +863,11 @@ impl Executable for RecursiveCanonical { let type_ids = type_ids.execute::(ctx)?.0.into_array(); let children = children .iter() + .cloned() .map(|child| { - Ok(child - .clone() - .execute::(ctx)? - .0 - .into_array()) + child + .execute::(ctx) + .map(|canonical| canonical.0.into_array()) }) .collect::>>()?; From 7ab189477aa45f5a2e037b58c23029aad1cdd29c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 16 Jul 2026 10:54:39 -0400 Subject: [PATCH 4/6] Implement nullable UnionArray semantics Signed-off-by: Connor Tsui --- .../src/arrays/chunked/vtable/canonical.rs | 4 +- .../src/arrays/constant/vtable/canonical.rs | 45 ++--- vortex-array/src/arrays/masked/execute.rs | 20 +- vortex-array/src/arrays/union/array.rs | 53 +++-- vortex-array/src/arrays/union/compute/mask.rs | 27 +++ vortex-array/src/arrays/union/compute/mod.rs | 2 + .../src/arrays/union/compute/rules.rs | 2 + vortex-array/src/arrays/union/mod.rs | 14 +- vortex-array/src/arrays/union/tests.rs | 189 ++++++++++++++---- vortex-array/src/arrays/union/vtable/mod.rs | 17 +- .../src/arrays/union/vtable/operations.rs | 17 +- .../src/arrays/union/vtable/validity.rs | 6 +- vortex-array/src/canonical.rs | 4 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 30 ++- 14 files changed, 311 insertions(+), 119 deletions(-) create mode 100644 vortex-array/src/arrays/union/compute/mask.rs diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index dcf8af39ae5..e0089043828 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -25,8 +25,8 @@ 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::TYPE_IDS_DTYPE; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::type_ids_dtype; use crate::arrays::variant::VariantArrayExt; use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; @@ -105,7 +105,7 @@ fn pack_union_chunks(chunks: Vec, ctx: &mut ExecutionCtx) -> VortexRes .iter() .map(|chunk| chunk.type_ids().clone()) .collect(), - TYPE_IDS_DTYPE, + type_ids_dtype(union_chunks[0].dtype().nullability()), )? .into_array(); let children = variants diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index ff4c1784700..818c46b8394 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -8,7 +8,6 @@ use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::Canonical; use crate::ExecutionCtx; @@ -32,6 +31,7 @@ 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; @@ -167,9 +167,12 @@ pub(crate) fn constant_canonicalize( StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) }) } - DType::Union(variants) => { - Canonical::Union(constant_canonical_union(scalar, variants, array.len())?) - } + 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, @@ -196,41 +199,31 @@ pub(crate) fn constant_canonicalize( fn constant_canonical_union( scalar: &Scalar, variants: &UnionVariants, + nullability: Nullability, len: usize, ) -> VortexResult { - if scalar.is_null() { - vortex_bail!( - "Canonicalizing a null union scalar is not supported until null semantics are defined" - ) - } - if variants.variants().any(|variant| variant.is_nullable()) { - vortex_bail!("Canonical UnionArray children must be non-nullable") - } - let union = scalar.as_union(); - let type_id = union - .type_id() - .vortex_expect("non-null union scalar must have a type ID"); - let selected_child = union - .child_index() - .vortex_expect("validated union scalar must select a child"); - let selected_value = union - .value() - .vortex_expect("non-null union scalar must have a value"); + let selected = union.child_index().zip(union.value()); let children = variants .variants() .enumerate() .map(|(index, dtype)| { - let value = if index == selected_child { - selected_value.clone() - } else { - Scalar::zero_value(&dtype) + let value = match selected { + Some((selected_child, selected_value)) if index == selected_child => { + selected_value.clone() + } + _ => Scalar::default_value(&dtype), }; 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(), diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index c5fc9ac90af..5c9be9e6370 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::Canonical; use crate::IntoArray; @@ -18,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; @@ -25,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; @@ -51,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(_) => { - vortex_bail!("Masking UnionArray is not supported until null semantics are defined") - } + 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)?), }) @@ -146,6 +145,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/union/array.rs b/vortex-array/src/arrays/union/array.rs index fef81d2672d..0c58f91f4f0 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -19,8 +19,8 @@ use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; use crate::arrays::PrimitiveArray; use crate::arrays::Union; -use crate::arrays::union::TYPE_IDS_DTYPE; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::dtype::UnionVariants; use crate::legacy_session; @@ -35,11 +35,18 @@ pub(super) fn make_union_parts( 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), len, EmptyArrayData).with_slots(slots) + ArrayParts::new( + Union, + DType::Union(variants, nullability), + len, + EmptyArrayData, + ) + .with_slots(slots) } /// Concrete parts of a [`UnionArray`](super::UnionArray). @@ -57,12 +64,12 @@ pub trait UnionArrayExt: TypedArrayRef { /// The union's variant schema. fn variants(&self) -> &UnionVariants { match self.as_ref().dtype() { - DType::Union(variants) => variants, + DType::Union(variants, _) => variants, _ => unreachable!("UnionArrayExt requires a union dtype"), } } - /// The row-aligned non-nullable `i8` type IDs. + /// 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() @@ -87,7 +94,7 @@ pub trait UnionArrayExt: TypedArrayRef { } /// Return a sparse child selected by a data-level type ID. - fn child_by_type_id(&self, type_id: i8) -> Option<&ArrayRef> { + fn child_by_type_id(&self, type_id: u8) -> Option<&ArrayRef> { self.child(self.variants().tag_to_child_index(type_id)?) } @@ -118,13 +125,22 @@ fn validate_type_id_values(array: &Array) -> VortexResult<()> { return Ok(()); } + let mut ctx = legacy_session().create_execution_ctx(); let type_ids = array .type_ids() .clone() - .execute::(&mut legacy_session().create_execution_ctx())?; - for type_id in type_ids.as_slice::() { + .execute::(&mut ctx)?; + let values = type_ids.as_slice::(); + let validity = type_ids.validity()?.execute_mask(array.len(), &mut ctx)?; + let valid_indices: Box + '_> = match &validity { + vortex_mask::Mask::AllTrue(_) => Box::new(0..values.len()), + vortex_mask::Mask::AllFalse(_) => Box::new(std::iter::empty()), + vortex_mask::Mask::Values(mask) => Box::new(mask.indices().iter().copied()), + }; + for index in valid_indices { + let type_id = values[index]; vortex_ensure!( - array.variants().tag_to_child_index(*type_id).is_some(), + array.variants().tag_to_child_index(type_id).is_some(), "UnionArray type ID {type_id} is not present in {:?}", array.variants().type_ids() ); @@ -150,8 +166,6 @@ impl Array { /// Try to construct a canonical sparse union array. /// - /// Until Union nullability semantics are settled, every child must be non-nullable. - /// /// # Errors /// /// Returns an error if the type IDs and sparse children do not match the variant schema or do @@ -162,8 +176,11 @@ impl Array { children: impl Into>, ) -> VortexResult { vortex_ensure!( - type_ids.dtype() == &TYPE_IDS_DTYPE, - "UnionArray type_ids must be non-nullable i8, got {}", + matches!( + type_ids.dtype(), + DType::Primitive(crate::dtype::PType::U8, _) + ), + "UnionArray type_ids must be u8, got {}", type_ids.dtype() ); @@ -178,9 +195,9 @@ impl Array { /// /// # Safety /// - /// The caller must ensure the type IDs are non-nullable `i8` values declared by `variants`, - /// every child has the corresponding variant dtype, and all arrays have the same length. The - /// children must currently be non-nullable. + /// The caller must ensure every non-null type ID is a `u8` value declared by `variants`, 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, @@ -202,9 +219,9 @@ impl Array { } } - /// Create an empty array for a non-nullable union dtype. - pub(crate) fn empty(variants: UnionVariants) -> Self { - let type_ids = PrimitiveArray::from_iter(Vec::::new()).into_array(); + /// 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()) 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 index e9208bf2185..8130c88ceb4 100644 --- a/vortex-array/src/arrays/union/compute/mod.rs +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -4,6 +4,8 @@ mod filter; pub(crate) use filter::filter_union; +mod mask; + pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/union/compute/rules.rs b/vortex-array/src/arrays/union/compute/rules.rs index 4666a99b92f..ce06e369cd4 100644 --- a/vortex-array/src/arrays/union/compute/rules.rs +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -6,9 +6,11 @@ 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/mod.rs b/vortex-array/src/arrays/union/mod.rs index b8671ccafd8..cce96577ed5 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -3,12 +3,10 @@ //! Canonical sparse union arrays. //! -//! A [`UnionArray`] stores one non-nullable `i8` 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. -//! -//! Union nullability semantics are still being designed, so this encoding currently requires every -//! variant child to be non-nullable. +//! 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. use crate::dtype::DType; use crate::dtype::Nullability; @@ -24,7 +22,9 @@ pub(crate) mod compute; mod vtable; pub use vtable::Union; -pub(crate) const TYPE_IDS_DTYPE: DType = DType::Primitive(PType::I8, Nullability::NonNullable); +pub(crate) fn 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 index e800647d56b..d8efe4a3c71 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests.rs @@ -19,6 +19,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; @@ -41,7 +42,7 @@ fn variants() -> VortexResult { fn union_array() -> VortexResult { UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 9, 5]).into_array(), + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), variants()?, vec![ PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), @@ -50,6 +51,28 @@ fn union_array() -> VortexResult { ) } +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()?; @@ -62,11 +85,11 @@ fn scalar_at_uses_type_id_indirection() -> VortexResult<()> { assert!(array.child_by_name_opt("missing").is_none()); assert_eq!( array.execute_scalar(0, &mut ctx)?, - Scalar::union(variants()?, 5, 10i32.into())? + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable,)? ); assert_eq!( array.execute_scalar(1, &mut ctx)?, - Scalar::union(variants()?, 9, true.into())? + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable,)? ); assert!(matches!(array.validity()?, Validity::NonNullable)); @@ -81,21 +104,21 @@ fn validates_sparse_components() -> VortexResult<()> { ]; let unknown_type_id = UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 7, 5]).into_array(), + PrimitiveArray::from_iter([5u8, 7, 5]).into_array(), variants()?, children.clone(), ); assert!(unknown_type_id.is_err()); let mismatched_lengths = UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 9]).into_array(), + PrimitiveArray::from_iter([5u8, 9]).into_array(), variants()?, children, ); assert!(mismatched_lengths.is_err()); let nullable_child = UnionArray::try_new( - PrimitiveArray::from_iter([5i8, 9, 5]).into_array(), + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), variants()?, vec![ PrimitiveArray::new(buffer![10i32, 0, 30], Validity::AllValid).into_array(), @@ -107,6 +130,69 @@ fn validates_sparse_components() -> VortexResult<()> { Ok(()) } +#[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(); @@ -118,15 +204,15 @@ fn structural_operations_preserve_sparse_alignment() -> VortexResult<()> { assert_eq!( sliced.execute_scalar(0, &mut ctx)?, - Scalar::union(variants()?, 9, true.into())? + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable,)? ); assert_eq!( filtered.execute_scalar(1, &mut ctx)?, - Scalar::union(variants()?, 5, 30i32.into())? + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? ); assert_eq!( taken.execute_scalar(0, &mut ctx)?, - Scalar::union(variants()?, 5, 30i32.into())? + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? ); Ok(()) @@ -136,33 +222,34 @@ fn structural_operations_preserve_sparse_alignment() -> VortexResult<()> { fn serde_roundtrip() -> VortexResult<()> { let session = array_session(); let mut execution_ctx = session.create_execution_ctx(); - let array = 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)? - ); + 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(()) @@ -170,7 +257,7 @@ fn serde_roundtrip() -> VortexResult<()> { #[test] fn constant_union_executes_to_sparse_union() -> VortexResult<()> { - let scalar = Scalar::union(variants()?, 9, true.into())?; + 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() @@ -179,6 +266,32 @@ fn constant_union_executes_to_sparse_union() -> VortexResult<()> { 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(()) } @@ -195,7 +308,7 @@ fn chunked_union_packs_components() -> VortexResult<()> { assert!(canonical.iter_children().all(|child| child.is::())); assert_eq!( canonical.execute_scalar(2, &mut ctx)?, - Scalar::union(variants()?, 5, 30i32.into())? + 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 index 4023699c256..2bd56b941c2 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -19,12 +19,12 @@ use crate::array::ArrayView; use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::union::TYPE_IDS_DTYPE; 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::type_ids_dtype; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -57,13 +57,9 @@ impl VTable for Union { len: usize, slots: &[Option], ) -> VortexResult<()> { - let DType::Union(variants) = dtype else { + let DType::Union(variants, nullability) = dtype else { vortex_bail!("Expected union dtype, found {dtype}") }; - vortex_ensure!( - variants.variants().all(|variant| !variant.is_nullable()), - "UnionArray children must be non-nullable" - ); vortex_ensure!( slots.len() == CHILDREN_OFFSET + variants.len(), "UnionArray has {} slots but expected {}", @@ -74,9 +70,10 @@ impl VTable for Union { let type_ids = slots[TYPE_IDS_SLOT] .as_ref() .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?; + let expected_type_ids_dtype = type_ids_dtype(*nullability); vortex_ensure!( - type_ids.dtype() == &TYPE_IDS_DTYPE, - "UnionArray type_ids must be non-nullable i8, got {}", + type_ids.dtype() == &expected_type_ids_dtype, + "UnionArray type_ids must have dtype {expected_type_ids_dtype}, got {}", type_ids.dtype() ); vortex_ensure!( @@ -142,7 +139,7 @@ impl VTable for Union { ) -> VortexResult> { vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); - let DType::Union(variants) = dtype else { + let DType::Union(variants, nullability) = dtype else { vortex_bail!("Expected union dtype, found {dtype}") }; vortex_ensure!( @@ -152,7 +149,7 @@ impl VTable for Union { children.len() ); - let type_ids = children.get(TYPE_IDS_SLOT, &TYPE_IDS_DTYPE, len)?; + let type_ids = children.get(TYPE_IDS_SLOT, &type_ids_dtype(*nullability), len)?; let sparse_children = variants .variants() .enumerate() diff --git a/vortex-array/src/arrays/union/vtable/operations.rs b/vortex-array/src/arrays/union/vtable/operations.rs index debd1b4ebac..3baef4c79f3 100644 --- a/vortex-array/src/arrays/union/vtable/operations.rs +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -17,12 +17,10 @@ impl OperationsVTable for Union { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let type_id = array - .type_ids() - .execute_scalar(index, ctx)? - .as_primitive() - .typed_value::() - .ok_or_else(|| vortex_err!("UnionArray type ID at index {index} is null"))?; + 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 child_index = array .variants() .tag_to_child_index(type_id) @@ -32,6 +30,11 @@ impl OperationsVTable for Union { .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) + Scalar::union( + array.variants().clone(), + type_id, + child_scalar, + array.dtype().nullability(), + ) } } diff --git a/vortex-array/src/arrays/union/vtable/validity.rs b/vortex-array/src/arrays/union/vtable/validity.rs index 0231b7a3334..80e043534e5 100644 --- a/vortex-array/src/arrays/union/vtable/validity.rs +++ b/vortex-array/src/arrays/union/vtable/validity.rs @@ -6,11 +6,11 @@ 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 { - // TODO(connor)[Union]: define how union and child nullability interact. - Ok(Validity::NonNullable) + 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 2e4851ac15d..9447d0831ae 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -234,7 +234,9 @@ impl Canonical { Validity::from(n), ) }), - DType::Union(variants) => Canonical::Union(UnionArray::empty(variants.clone())), + DType::Union(variants, nullability) => { + Canonical::Union(UnionArray::empty(variants.clone(), *nullability)) + } DType::Variant(_) => { vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant") } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 1959da25a1d..702aadbea7b 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::type_ids_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::expression::Expression; @@ -185,9 +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(_) => { - vortex_bail!("Union arrays don't support casting yet") - } + CanonicalView::Union(a) => cast_union(a, dtype), CanonicalView::Extension(a) => ::cast(a, dtype), CanonicalView::Variant(_) => { vortex_bail!("Variant arrays don't support casting") @@ -195,6 +197,28 @@ 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(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) From e7465e5cea0c1d13b8c27476ca1fd1b794fa352c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 16 Jul 2026 11:45:58 -0400 Subject: [PATCH 5/6] Handle nested Union constant placeholders Signed-off-by: Connor Tsui --- .../src/arrays/constant/vtable/canonical.rs | 3 ++ vortex-array/src/arrays/union/tests.rs | 44 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 818c46b8394..6ade84a2ac1 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -213,6 +213,9 @@ fn constant_canonical_union( Some((selected_child, selected_value)) if index == selected_child => { selected_value.clone() } + _ if matches!(&dtype, DType::Union(_, Nullability::NonNullable)) => { + Scalar::zero_value(&dtype) + } _ => Scalar::default_value(&dtype), }; ConstantArray::new(value, len).into_array() diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests.rs index d8efe4a3c71..7a3ce1ce631 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests.rs @@ -295,6 +295,50 @@ fn constant_union_executes_to_sparse_union() -> VortexResult<()> { 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_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.clone(), + ], + 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)?, + Scalar::zero_value(&nested_dtype) + ); + + 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)?, + Scalar::zero_value(&nested_dtype) + ); + + Ok(()) +} + #[test] fn chunked_union_packs_components() -> VortexResult<()> { let first = union_array()?.into_array().slice(0..1)?; From 16dddab1c0dbb03f7bb9687270e7c12dab8f888f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 16 Jul 2026 14:47:40 -0400 Subject: [PATCH 6/6] Fix UnionArray canonicalization edge cases Signed-off-by: Connor Tsui --- .../src/arrays/constant/vtable/canonical.rs | 38 ++++-- vortex-array/src/arrays/union/tests.rs | 118 ++++++++++++++++++ vortex-array/src/canonical.rs | 26 +++- 3 files changed, 166 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 6ade84a2ac1..010b09d2c96 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; @@ -149,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 = Scalar::try_default_value(&dt).ok_or_else(|| { + vortex_err!( + "cannot canonicalize null constant struct: field dtype {dt} \ + has no default value" + ) + })?; + Ok(ConstantArray::new(scalar, array.len()).into_array()) }) - .collect() + .collect::>>()? } }; // SAFETY: Fields are constructed from the same struct scalar, all have same @@ -202,6 +209,10 @@ fn constant_canonical_union( 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.value()); @@ -213,14 +224,17 @@ fn constant_canonical_union( Some((selected_child, selected_value)) if index == selected_child => { selected_value.clone() } - _ if matches!(&dtype, DType::Union(_, Nullability::NonNullable)) => { - Scalar::zero_value(&dtype) - } - _ => Scalar::default_value(&dtype), + _ => Scalar::try_default_value(&dtype).ok_or_else(|| { + vortex_err!( + "cannot canonicalize constant union: unselected variant {} ({dtype}) has \ + no default value", + variants.names()[index] + ) + })?, }; - ConstantArray::new(value, len).into_array() + Ok(ConstantArray::new(value, len).into_array()) }) - .collect::>(); + .collect::>>()?; let type_id = match union.type_id() { Some(type_id) => Scalar::primitive(type_id, nullability), diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests.rs index 7a3ce1ce631..92e47b69c9b 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests.rs @@ -4,10 +4,13 @@ use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; 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; @@ -339,6 +342,121 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> { Ok(()) } +#[test] +fn constant_union_builds_deeply_nested_union_placeholders() -> VortexResult<()> { + let nested_dtype = DType::Union( + UnionVariants::try_new( + ["value"].into(), + vec![DType::Primitive(PType::I64, Nullability::NonNullable)], + vec![3], + )?, + Nullability::NonNullable, + ); + let wrapper_dtype = + DType::struct_([("nested", nested_dtype.clone())], 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(Scalar::default_value(&nested_dtype)) + ); + + Ok(()) +} + +#[test] +fn constant_union_rejects_uninhabited_placeholders_without_panicking() -> VortexResult<()> { + let uninhabited = DType::Union( + UnionVariants::new(Default::default(), vec![])?, + Nullability::NonNullable, + ); + let outer_variants = UnionVariants::try_new( + ["number", "uninhabited"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + uninhabited, + ], + vec![5, 9], + )?; + let selected_number = + Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?; + let mut ctx = array_session().create_execution_ctx(); + + assert!( + ConstantArray::new(selected_number.clone(), 1) + .into_array() + .execute::(&mut ctx) + .is_err() + ); + assert_eq!( + ConstantArray::new(selected_number, 0) + .into_array() + .execute::(&mut ctx)? + .len(), + 0 + ); + + 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)?; diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 9447d0831ae..f01a7cac41c 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; @@ -237,9 +238,15 @@ impl Canonical { DType::Union(variants, nullability) => { Canonical::Union(UnionArray::empty(variants.clone(), *nullability)) } - DType::Variant(_) => { - vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant") - } + 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(), @@ -678,7 +685,18 @@ impl Executable for CanonicalValidity { StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?) }))) } - union @ Canonical::Union(_) => Ok(CanonicalValidity(union)), + 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(),