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..76cbd8f36d7 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,48 @@ impl Scalar { .vortex_expect("unable to construct an extension `Scalar`") } + /// Creates a union scalar from a type ID and child scalar. + /// + /// The selected child's dtype is verified and then discarded. Its raw value is stored 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.into_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..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..2d36ca3e0f7 100644 --- a/vortex-array/src/scalar/downcast.rs +++ b/vortex-array/src/scalar/downcast.rs @@ -19,6 +19,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 +158,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 !self.dtype().is_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 +295,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..362d08e718c 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(ScalarValue::to_proto(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_value = ScalarValue::from_proto(child_proto, &child_dtype, session)?; + Scalar::validate(&child_dtype, child_value.as_ref()).map_err(|error| { + vortex_err!( + Serde: "union type ID {type_id} has invalid child for dtype {child_dtype}: {error}" + ) + })?; + + Ok(ScalarValue::Union(UnionValue::new(type_id, child_value))) +} + #[cfg(test)] mod tests { use std::f32; @@ -468,6 +517,7 @@ mod tests { use std::sync::Arc; use vortex_buffer::BufferString; + use vortex_error::VortexError; use vortex_error::vortex_panic; use vortex_proto::scalar as pb; use vortex_session::VortexSession; @@ -477,6 +527,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 +714,112 @@ 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(union_value)) + if matches!( + union_value + .value + .as_deref() + .and_then(|value| value.kind.as_ref()), + Some(Kind::NullValue(_)) + ) + )); + 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(_)) + )); + + assert_ne!(inner_null_proto, outer_null_proto); + round_trip(inner_null); + round_trip(outer_null); + + Ok(()) + } + + #[test] + fn test_union_proto_rejects_malformed_values() -> 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(ScalarValue::to_proto( + Scalar::primitive(42_i32, Nullability::Nullable).value(), + ))), + }))), + }; + + assert!(ScalarValue::from_proto(&unknown_tag, &dtype, &session()).is_err()); + + let missing_child = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 5, + value: None, + }))), + }; + + assert!(matches!( + ScalarValue::from_proto(&missing_child, &dtype, &session()), + Err(VortexError::Serde(..)) + )); + + let wrong_child_value = pb::ScalarValue { + kind: Some(Kind::UnionValue(Box::new(PbUnionValue { + type_id: 5, + value: Some(Box::new(ScalarValue::to_proto( + Scalar::utf8("wrong", Nullability::NonNullable).value(), + ))), + }))), + }; + + assert!(matches!( + ScalarValue::from_proto(&wrong_child_value, &dtype, &session()), + Err(VortexError::Serde(..)) + )); + + 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..92205aa81d6 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -100,21 +100,27 @@ 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. + /// + /// # Panics + /// + /// Panics if `dtype` has no default value, such as a 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` has no default. + 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 +133,12 @@ impl Scalar { /// element [`DType`] /// - `Struct`: A struct where each field has a zero value, which is determined by the field /// [`DType`] + /// - `Union`: Does not have a "zero" value /// - `Extension`: The zero value of the storage [`DType`] + /// + /// # Panics + /// + /// Panics if the dtype has no non-null zero value, such as `Null` or a union. pub fn zero_value(dtype: &DType) -> Self { let value = ScalarValue::zero_value(dtype); @@ -189,7 +200,7 @@ impl Scalar { DType::Binary(_) => value.as_binary().is_empty(), DType::List(..) => value.as_list().is_empty(), // A fixed-size list is zero only if it has the expected number of elements and every - // element is itself a non-null zero value. + // element is itself a non-null zero value.1 DType::FixedSizeList(_, list_size, _) => { let list = self.as_list(); list.len() == *list_size as usize @@ -201,7 +212,8 @@ 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"), + // We do not define union to have any valid "zero" value. + DType::Union(..) => false, 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,22 @@ 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_none() && rhs.value().is_none() { + return Some(Ordering::Equal); + } + 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_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, diff --git a/vortex-array/src/scalar/scalar_value.rs b/vortex-array/src/scalar/scalar_value.rs index ebc2101fac6..f042573f088 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 raw 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,40 +61,40 @@ 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(..) => return None, 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 has no default. /// - /// 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 a dtype with no default 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 +102,29 @@ 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(..) => return None, 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 +172,14 @@ impl Display for ScalarValue { } write!(f, "]") } + ScalarValue::Union(value) => { + write!(f, "union@{}(", value.type_id())?; + match value.value() { + Some(value) => write!(f, "{value}"), + None => write!(f, "null"), + }?; + write!(f, ")") + } 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..a21e7193497 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,238 @@ mod tests { PValue::F16(f16_value) ); } + + #[test] + fn union_casts_apply_inner_nullability() -> VortexResult<()> { + let source_variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + 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 unchanged_child = Scalar::union( + source_variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + let changed_child = Scalar::union( + source_variants, + 9, + Scalar::utf8("hello", Nullability::NonNullable), + Nullability::NonNullable, + )?; + + let unchanged_child = unchanged_child.cast(&target_dtype)?; + let changed_child = changed_child.cast(&target_dtype)?; + + assert_eq!(unchanged_child.dtype(), &target_dtype); + assert_eq!( + unchanged_child + .as_union() + .value() + .map(|scalar| scalar.dtype().clone()), + Some(DType::Primitive(PType::I32, Nullability::NonNullable)) + ); + assert_eq!(changed_child.dtype(), &target_dtype); + assert_eq!( + changed_child + .as_union() + .value() + .map(|scalar| scalar.dtype().clone()), + Some(DType::Utf8(Nullability::Nullable)) + ); + + Ok(()) + } + + #[test] + fn union_cast_rejects_null_for_non_nullable_child() -> VortexResult<()> { + let source_variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let target_variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![5], + )?; + let scalar = Scalar::union( + source_variants, + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::NonNullable, + )?; + + assert!( + scalar + .cast(&DType::Union(target_variants, Nullability::NonNullable)) + .is_err() + ); + + Ok(()) + } + + #[test] + fn union_cast_widens_outer_nullability_with_selected_inner_null() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let scalar = Scalar::union( + variants.clone(), + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::NonNullable, + )?; + let target_dtype = DType::Union(variants, Nullability::Nullable); + + let cast = scalar.cast(&target_dtype)?; + let selected_child = cast + .as_union() + .value() + .vortex_expect("present union must have a selected child"); + + assert_eq!(cast.dtype(), &target_dtype); + assert!(selected_child.is_null()); + assert_eq!( + selected_child.dtype(), + &DType::Primitive(PType::I32, Nullability::Nullable) + ); + + Ok(()) + } + + #[test] + fn union_nullability_cast_recurses_inside_list() -> VortexResult<()> { + let source_variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![5], + )?; + let source_union_dtype = DType::Union(source_variants.clone(), Nullability::NonNullable); + let scalar = Scalar::list( + source_union_dtype, + vec![Scalar::union( + source_variants, + 5, + Scalar::primitive(42_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?], + Nullability::NonNullable, + ); + + let target_variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let target_union_dtype = DType::Union(target_variants, Nullability::NonNullable); + let target_dtype = DType::List(Arc::new(target_union_dtype.clone()), Nullability::Nullable); + + let cast = scalar.cast(&target_dtype)?; + let union = cast + .as_list() + .element(0) + .vortex_expect("list must contain its union child"); + + assert_eq!(cast.dtype(), &target_dtype); + assert_eq!(union.dtype(), &target_union_dtype); + assert_eq!( + union.as_union().value(), + Some(Scalar::primitive(42_i32, Nullability::Nullable)) + ); + + Ok(()) + } + + #[test] + fn union_cast_rejects_schema_changes() -> VortexResult<()> { + let source_variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + vec![5], + )?; + let scalar = Scalar::union( + source_variants, + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::NonNullable, + )?; + + let changed_name = UnionVariants::try_new( + ["integer"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![5], + )?; + let changed_type_id = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![6], + )?; + let changed_dtype = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I64, Nullability::Nullable)], + vec![5], + )?; + + for variants in [changed_name, changed_type_id, changed_dtype] { + assert!( + scalar + .cast(&DType::Union(variants, Nullability::NonNullable)) + .is_err() + ); + } + + Ok(()) + } + + #[test] + fn cast_changes_only_outer_union_nullability() -> 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 expected = DType::Union(variants, Nullability::Nullable); + let nullable = scalar.cast(&expected)?; + + 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| scalar.dtype().clone()), + 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..365d0dddd5f 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,60 @@ 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] + #[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"), + Nullability::NonNullable, + ); + + Scalar::default_value(&non_nullable); + } + + #[test] + #[should_panic(expected = "has no non-null zero value")] + fn union_has_no_zero_value() { + let dtype = DType::Union( + union_variants(Nullability::NonNullable, Nullability::NonNullable) + .vortex_expect("union variants must be valid"), + Nullability::NonNullable, + ); + + Scalar::zero_value(&dtype); + } + + #[test] + fn union_is_zero_does_not_inspect_selected_child() -> VortexResult<()> { + let variants = union_variants(Nullability::NonNullable, Nullability::NonNullable)?; + let scalar = Scalar::union( + variants.clone(), + 5, + Scalar::from(0_i32), + Nullability::NonNullable, + )?; + let null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + + assert_eq!(scalar.is_zero(), Some(false)); + assert_eq!(null.is_zero(), None); + + Ok(()) + } + + #[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 +212,36 @@ 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::() + ); + let inner_null = Scalar::union( + variants.clone(), + 5, + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable, + )?; + assert_eq!( + inner_null.approx_nbytes(), + size_of::() + size_of::() + ); + assert_eq!( + Scalar::null(DType::Union(variants, Nullability::Nullable)).approx_nbytes(), + 0 + ); + + Ok(()) } #[test] @@ -438,6 +538,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 +606,100 @@ 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 union_values_with_same_type_id_are_ordered_by_selected_child() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![3], + )?; + let smaller = Scalar::union( + variants.clone(), + 3, + Scalar::primitive(10_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + let larger = Scalar::union( + variants, + 3, + Scalar::primitive(20_i32, Nullability::NonNullable), + Nullability::NonNullable, + )?; + + assert_eq!(smaller.partial_cmp(&larger), Some(std::cmp::Ordering::Less)); + assert_eq!( + larger.partial_cmp(&smaller), + Some(std::cmp::Ordering::Greater) + ); + + 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.clone(), + 8, + Scalar::null(DType::Utf8(Nullability::Nullable)), + Nullability::NonNullable, + )?; + let outer_null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + + 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)); + assert_ne!(int_null, outer_null); + assert_ne!(scalar_hash(&int_null), scalar_hash(&outer_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..30fe6e00b76 --- /dev/null +++ b/vortex-array/src/scalar/typed_view/union.rs @@ -0,0 +1,353 @@ +// 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_ensure; +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 selected +/// child's raw value is optional so a present union whose selected child is null remains distinct +/// from an outer null union. Its dtype is determined by `type_id` and the enclosing +/// [`DType::Union`]. +#[derive(Debug, Clone)] +pub struct UnionValue { + /// The type ID selecting a variant in the enclosing [`DType::Union`]. + type_id: u8, + /// The selected variant's raw value, or [`None`] if the selected child is null. + /// + /// This is boxed to break the recursive layout between [`Scalar`], [`ScalarValue`], and + /// [`UnionValue`]. + value: Option>, +} + +impl PartialEq for UnionValue { + fn eq(&self, other: &Self) -> bool { + if self.value.is_none() && other.value.is_none() { + 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_none(); + 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: Option) -> Self { + Self { + type_id, + value: value.map(Box::new), + } + } + + /// Returns the type ID selecting the union variant. + #[inline] + pub fn type_id(&self) -> u8 { + self.type_id + } + + /// Returns the selected variant's raw value, or [`None`] if the selected child is null. + #[inline] + pub fn value(&self) -> Option<&ScalarValue> { + self.value.as_deref() + } +} + +/// A typed view into a [`DType::Union`] scalar. +/// +/// A present union scalar carries a type ID and the raw value 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"); + }; + + write!(f, "{name}(")?; + match self + .value + .vortex_expect("non-null union scalar must have a selected value") + .value() + { + Some(value) => write!(f, "{value}"), + None => write!(f, "null"), + }?; + write!(f, ")") + } +} + +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 { + vortex_ensure!(dtype.is_union(), "Expected union scalar, found {dtype}"); + + 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 { + if !dtype.is_union() { + 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. + /// + /// A selected null child is returned as a present, null [`Scalar`]. + pub fn value(&self) -> Option { + let union_value = self.value?; + let child_dtype = self + .variant_dtype() + .vortex_expect("validated union type ID must select a child dtype"); + + // SAFETY: Construction of this `UnionScalar` guarantees that its type ID resolves to this + // child dtype and that the raw child value recursively validates against it. + Some(unsafe { Scalar::new_unchecked(child_dtype, union_value.value().cloned()) }) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexExpect; + 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)); + assert_eq!( + scalar + .value() + .vortex_expect("union must have an outer value") + .as_union() + .value(), + Some(&ScalarValue::Primitive(42_i32.into())) + ); + + 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| scalar.is_null())); + assert!( + scalar + .value() + .vortex_expect("union must have an outer value") + .as_union() + .value() + .is_none() + ); + + 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_rejects_non_union_dtype() { + assert!( + UnionScalar::try_new(&DType::Bool(Nullability::Nullable), None).is_err(), + "non-union dtypes must return an error" + ); + } + + #[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).into_value(), + )); + + assert!(UnionScalar::try_new(&dtype, Some(&unknown_type_id)).is_err()); + + let wrong_value = ScalarValue::Union(UnionValue::new( + 5, + Scalar::utf8("wrong", Nullability::NonNullable).into_value(), + )); + + 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..45d9677f361 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"); + + Self::validate(&child_dtype, union_value.value()).map_err(|error| { + vortex_error::vortex_err!( + "union value for type ID {type_id} is invalid for dtype {child_dtype}: \ + {error}" + ) + })?; + } 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).into_value(), + ))), + ) + .is_err() + ); + + assert!( + Scalar::try_new( + dtype, + Some(ScalarValue::Union(UnionValue::new( + 5, + Scalar::utf8("wrong", Nullability::NonNullable).into_value(), + ))), + ) + .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..90e9afe6a6e 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; + ScalarValue value = 2; +} diff --git a/vortex-proto/src/generated/vortex.scalar.rs b/vortex-proto/src/generated/vortex.scalar.rs index df43794acf7..7ea4e529104 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>, +}