diff --git a/vortex-array/src/arrays/constant/vtable/operations.rs b/vortex-array/src/arrays/constant/vtable/operations.rs index fdd2dd67345..8162a86939a 100644 --- a/vortex-array/src/arrays/constant/vtable/operations.rs +++ b/vortex-array/src/arrays/constant/vtable/operations.rs @@ -18,3 +18,46 @@ 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), + )?; + 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)); + 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..17265297641 100644 --- a/vortex-array/src/scalar/arbitrary.rs +++ b/vortex-array/src/scalar/arbitrary.rs @@ -95,7 +95,21 @@ 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) => { + 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, + ) + .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..1fb57a99511 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -24,6 +24,16 @@ impl Scalar { return Ok(self.clone()); } + // Union nullability is part of its variant dtypes, so the generic nullability-only cast + // below is not valid for unions. Keep all non-identity union casts unsupported until their + // semantics are defined. + if self.dtype().is_union() || target_dtype.is_union() { + vortex_bail!( + "non-identity union scalar cast from {} to {target_dtype} is not supported", + self.dtype() + ); + } + // Check for solely nullability casting. if self.dtype().eq_ignore_nullability(target_dtype) { // Cast from non-nullable to nullable or vice versa. @@ -58,13 +68,16 @@ 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(..) => unreachable!("union casts are handled before scalar dispatch"), DType::Variant(_) => vortex_bail!("Variant scalars can't be cast to {target_dtype}"), DType::Extension(..) => self.as_extension().cast(target_dtype), } } /// Cast the scalar into a nullable version of its current type. + /// + /// For a union this makes every variant nullable, since a union has no top-level nullability of + /// its own. pub fn into_nullable(self) -> Scalar { let (dtype, value) = self.into_parts(); Self::try_new(dtype.as_nullable(), value) diff --git a/vortex-array/src/scalar/constructor.rs b/vortex-array/src/scalar/constructor.rs index 16c87ca27a6..30e5aff0240 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,42 @@ impl Scalar { .vortex_expect("unable to construct an extension `Scalar`") } + /// Creates a union scalar from a type ID and child scalar. + /// + /// A null child is normalized to a null union scalar, so null union scalars do not retain a + /// selected type ID. A type ID is not part of a null union scalar's logical identity and is not + /// preserved across serialization or other round trips. The type ID and exact child dtype are + /// still validated before normalization. + /// + /// # 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: i8, child: Scalar) -> 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() + ); + + let (_, child_value) = child.into_parts(); + let value = child_value.map(|value| ScalarValue::Union(UnionValue::new(type_id, value))); + + Self::try_new(DType::Union(variants), value) + } + /// 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..19b7ad74863 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,24 @@ 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))?; + + assert_eq!(format!("{scalar}"), "int(42i32)"); + assert_eq!(format!("{}", Scalar::null(DType::Union(variants))), "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..463ea92251a 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..66255f8afe8 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: i32::from(v.type_id()), + value: Some(Box::new(ScalarValue::to_proto(Some(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 non-null 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 = i8::try_from(value.type_id).map_err( + |_| vortex_err!(Serde: "union type ID {} is outside the i8 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_value = + ScalarValue::from_proto(child_proto, &child_dtype, session)?.ok_or_else(|| { + vortex_err!( + Serde: "union type ID {type_id} has a nested null; null unions must use NullValue" + ) + })?; + + Ok(ScalarValue::Union(UnionValue::new(type_id, child_value))) +} + #[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,76 @@ 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), + )?); + + let null = Scalar::union( + variants, + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + )?; + let null_proto = pb::Scalar::from(&null); + + assert!(matches!( + null_proto + .value + .as_deref() + .and_then(|value| value.kind.as_ref()), + Some(Kind::NullValue(_)) + )); + round_trip(null); + + Ok(()) + } + + #[test] + fn test_union_proto_rejects_nested_null_and_unknown_tag() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let dtype = DType::Union(variants); + + let nested_null = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 5, + value: Some(Box::new(pb::ScalarValue { + kind: Some(Kind::NullValue(0)), + })), + }))), + }; + + assert!(ScalarValue::from_proto(&nested_null, &dtype, &session()).is_err()); + + let unknown_tag = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 7, + value: Some(Box::new(pb::ScalarValue { + kind: Some(Kind::Int64Value(42)), + })), + }))), + }; + + assert!(ScalarValue::from_proto(&unknown_tag, &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 60a675e42b9..4f4034979b3 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -103,6 +103,10 @@ impl Scalar { /// 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. + /// + /// # Panics + /// + /// Panics if `dtype` is a non-nullable union, which has no default value. pub fn default_value(dtype: &DType) -> Self { let value = ScalarValue::default_value(dtype); @@ -114,7 +118,7 @@ impl Scalar { /// /// # 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 +131,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 +210,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,18 +279,20 @@ 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(), } } } -/// We implement `Hash` manually to be consistent with `PartialEq`. Since we ignore nullability in -/// equality comparisons, we must also ignore it when hashing to maintain the invariant that equal -/// values have equal hashes. +/// We implement `Hash` manually so top-level dtype nullability does not affect the hash. impl Hash for Scalar { fn hash(&self, state: &mut H) { + // TODO(connor): Recursively ignore dtype nullability to match Scalar equality (#8744). self.dtype.as_nonnullable().hash(state); self.value.hash(state); } @@ -377,6 +388,20 @@ 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.type_id() != rhs.type_id() { + return None; + } + + let DType::Union(variants) = dtype else { + return None; + }; + + let child_index = variants.tag_to_child_index(lhs.type_id())?; + let child_dtype = variants.variant_by_index(child_index)?; + + partial_cmp_non_null_scalar_values(&child_dtype, lhs.value(), 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)); + + 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..f552546ece7 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 non-null union value carrying its selected type ID and child value. + 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,25 +61,34 @@ 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, value) = variants + .variants() + .enumerate() + .find_map(|(index, dtype)| Self::try_zero_value(&dtype).map(|v| (index, v)))?; + + Self::Union(UnionValue::new(variants.child_index_to_tag(index), value)) + } 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 @@ -97,10 +115,13 @@ impl ScalarValue { 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::default_value(&field)) + .collect(); Self::Tuple(field_values) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(_) => vortex_panic!("{dtype} has no default 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 @@ -156,6 +177,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..2a03f042906 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,77 @@ 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); + + let value = Scalar::union( + source_variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + )?; + + assert!(value.cast(&target_dtype).is_err()); + + let null = Scalar::null(DType::Union(source_variants)); + + assert!(null.cast(&target_dtype).is_err()); + + Ok(()) + } + + #[test] + fn into_nullable_makes_every_variant_nullable() -> VortexResult<()> { + let scalar = Scalar::union( + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?, + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + )?; + + let nullable = scalar.into_nullable(); + + // A union has no top-level nullability, so `into_nullable` propagates into every variant. + let expected = DType::Union(UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + ], + vec![5, 9], + )?); + 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(), + Some(Scalar::primitive(42_i32, Nullability::Nullable)) + ); + + Ok(()) + } } diff --git a/vortex-array/src/scalar/tests/primitives.rs b/vortex-array/src/scalar/tests/primitives.rs index d2e32ce6ef0..13d7183cc1a 100644 --- a/vortex-array/src/scalar/tests/primitives.rs +++ b/vortex-array/src/scalar/tests/primitives.rs @@ -8,6 +8,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; @@ -15,6 +17,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; @@ -24,6 +27,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], + ) + } + #[test] fn default_value_for_complex_dtype() { let struct_dtype = DType::struct_( @@ -57,7 +74,30 @@ 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, + )?); + + assert!(Scalar::default_value(&nullable).is_null()); + + Ok(()) + } + + #[test] + #[should_panic(expected = "has no default value")] + fn default_value_for_non_nullable_union_panics() { + let non_nullable = DType::Union( + union_variants(Nullability::NonNullable, Nullability::NonNullable) + .vortex_expect("union variants must be valid"), + ); + + Scalar::default_value(&non_nullable); + } + + #[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); @@ -133,6 +173,22 @@ 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), + )?; + + assert_eq!( + union_scalar.approx_nbytes(), + size_of::() + size_of::() + ); + assert_eq!(Scalar::null(DType::Union(variants)).approx_nbytes(), 0); + + Ok(()) } #[test] @@ -412,6 +468,33 @@ mod tests { assert_eq!(set.len(), 5); } + #[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), + )?; + + let rhs = Scalar::union( + rhs_variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + )?; + + assert_eq!(lhs, rhs); + + let lhs_null = Scalar::null(DType::Union(lhs_variants)); + let rhs_null = Scalar::null(DType::Union(rhs_variants)); + + assert_eq!(lhs_null, rhs_null); + + Ok(()) + } + #[test] fn test_scalar_partial_ord_incompatible_types() { let int_scalar = Scalar::primitive(42i32, Nullability::NonNullable); @@ -451,4 +534,32 @@ 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), + )?; + + let right = Scalar::union( + variants, + 8, + Scalar::primitive(42_i32, Nullability::NonNullable), + )?; + + assert_ne!(left, right); + assert_eq!(left.partial_cmp(&right), None); + + 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..59ff893eaad --- /dev/null +++ b/vortex-array/src/scalar/typed_view/union.rs @@ -0,0 +1,276 @@ +// 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 vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; + +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::UnionVariants; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; + +/// The non-null value stored by a union scalar. +/// +/// A null union is represented by the enclosing [`Scalar`]'s value being `None`, rather than by a +/// nested null in this type. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct UnionValue { + /// The type ID selecting a variant in the enclosing [`DType::Union`]. + type_id: i8, + /// The selected variant's non-null value, which must be valid for its dtype. + /// + /// This is boxed to break the recursive layout between [`ScalarValue`] and [`UnionValue`]. + value: Box, +} + +impl UnionValue { + pub(crate) fn new(type_id: i8, value: ScalarValue) -> Self { + Self { + type_id, + value: Box::new(value), + } + } + + /// Returns the type ID selecting the union variant. + #[inline] + pub fn type_id(&self) -> i8 { + self.type_id + } + + /// Returns the selected variant's non-null value. + #[inline] + pub fn value(&self) -> &ScalarValue { + &self.value + } +} + +/// A typed view into a [`DType::Union`] scalar. +/// +/// A non-null union scalar carries a type ID and a non-null value of the selected variant. A null +/// union scalar has neither because nullness is represented by the enclosing [`Scalar`]. +/// +/// The type ID is therefore not preserved across round trips for null union scalars. +#[derive(Debug, Clone, Copy)] +pub struct UnionScalar<'a> { + /// The variants of the union. + /// + /// The enclosing scalar's dtype is always `DType::Union(variants)`, so we hold the unwrapped + /// variants directly rather than re-matching the dtype on every access. + variants: &'a UnionVariants, + /// 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 { + // Match the nullability-agnostic equality of `Scalar`: compare the variant schema ignoring + // per-variant nullability, then the selected value. + self.variants.names() == other.variants.names() + && self.variants.type_ids() == other.variants.type_ids() + && self + .variants + .variants() + .zip(other.variants.variants()) + .all(|(lhs, rhs)| lhs.eq_ignore_nullability(&rhs)) + && 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(variants) = dtype else { + vortex_panic!("Expected union scalar, found {dtype}") + }; + + Self { + variants, + value: value.map(ScalarValue::as_union), + } + } + + /// Returns the variants of this union scalar. + #[inline] + pub fn variants(&self) -> &'a UnionVariants { + self.variants + } + + /// 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()?) + } + + /// Reconstructs the selected child scalar, or returns `None` if this scalar is null. + pub fn value(&self) -> Option { + let value = self.value?; + let dtype = self.variant_dtype()?; + + // SAFETY: Union scalar validation checks the type ID and validates this value against the + // selected variant dtype. + Some(unsafe { Scalar::new_unchecked(dtype, Some(value.value().clone())) }) + } +} + +#[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 scalar = Scalar::union( + variants()?, + 5, + Scalar::primitive(42_i32, 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.value(), + Some(Scalar::primitive(42_i32, Nullability::Nullable)) + ); + + Ok(()) + } + + #[test] + fn null_child_is_normalized_to_outer_null() -> VortexResult<()> { + let scalar = Scalar::union( + variants()?, + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + )?; + + let union = scalar.as_union(); + assert!(union.is_null()); + assert_eq!(union.type_id(), None); + assert_eq!(union.value(), None); + + Ok(()) + } + + #[test] + fn null_child_is_validated_before_normalization() -> VortexResult<()> { + let variants = variants()?; + let child = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); + + assert!(Scalar::union(variants.clone(), 7, child).is_err()); + + let wrong_child = Scalar::null(DType::Utf8(Nullability::Nullable)); + + assert!(Scalar::union(variants, 5, wrong_child).is_err()); + + Ok(()) + } + + #[test] + fn try_new_validates_type_id_and_selected_value() -> VortexResult<()> { + let dtype = DType::Union(variants()?); + let non_nullable_dtype = DType::Union(UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![5], + )?); + + assert!(UnionScalar::try_new(&non_nullable_dtype, None).is_err()); + + let unknown_type_id = + ScalarValue::Union(UnionValue::new(7, ScalarValue::Primitive(42_i32.into()))); + + assert!(UnionScalar::try_new(&dtype, Some(&unknown_type_id)).is_err()); + + let wrong_value = ScalarValue::Union(UnionValue::new(5, ScalarValue::Utf8("wrong".into()))); + + 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..eb1466cb609 100644 --- a/vortex-array/src/scalar/validate.rs +++ b/vortex-array/src/scalar/validate.rs @@ -1,10 +1,12 @@ // 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; use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use crate::dtype::DType; use crate::dtype::PType; @@ -118,7 +120,29 @@ 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"); + + Self::validate(&child_dtype, Some(union_value.value())).map_err(|err| { + vortex_err!( + "union value for type ID {type_id} is invalid for {child_dtype}: {err}" + ) + })?; + } 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); + + assert!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Union(UnionValue::new( + 7, + ScalarValue::Primitive(42_i32.into()), + ))), + ) + .is_err() + ); + + assert!( + Scalar::try_new( + dtype, + Some(ScalarValue::Union(UnionValue::new( + 5, + ScalarValue::Utf8("wrong".into()), + ))), + ) + .is_err() + ); + + Ok(()) + } +} diff --git a/vortex-proto/proto/scalar.proto b/vortex-proto/proto/scalar.proto index 251863dc3a3..c3f63b5f3cf 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 non-null union value. Null unions use ScalarValue.null_value instead. +message UnionValue { + int32 type_id = 1; + ScalarValue value = 2; +} diff --git a/vortex-proto/src/generated/vortex.scalar.rs b/vortex-proto/src/generated/vortex.scalar.rs index df43794acf7..9273f948d6c 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 non-null union value. Null unions use ScalarValue.null_value instead. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UnionValue { + #[prost(int32, tag = "1")] + pub type_id: i32, + #[prost(message, optional, boxed, tag = "2")] + pub value: ::core::option::Option<::prost::alloc::boxed::Box>, +}