diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 5767a124667..e3594f818c6 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -82,8 +82,9 @@ impl DType { /// Get a new DType with the given nullability (but otherwise the same as `self`). /// - /// [`DType::Null`] and [`DType::Union`] have intrinsic nullability and are returned unchanged. - /// To change a union's nullability, construct different [`UnionVariants`]. + /// [`DType::Null`] has intrinsic nullability and is returned unchanged. A [`DType::Union`] has + /// no top-level nullability of its own — its nullability is derived from its variants — so + /// `nullability` is propagated into every variant instead. pub fn with_nullability(&self, nullability: Nullability) -> Self { match self { Null => Null, @@ -95,7 +96,7 @@ impl DType { List(edt, _) => List(Arc::clone(edt), nullability), FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability), Struct(sf, _) => Struct(sf.clone(), nullability), - Union(vs) => Union(vs.clone()), + Union(vs) => Union(vs.with_nullability(nullability)), Variant(_) => Variant(nullability), Extension(ext) => Extension(ext.with_nullability(nullability)), } diff --git a/vortex-array/src/dtype/union.rs b/vortex-array/src/dtype/union.rs index faec75751c5..f2cd2ea84df 100644 --- a/vortex-array/src/dtype/union.rs +++ b/vortex-array/src/dtype/union.rs @@ -353,6 +353,21 @@ impl UnionVariants { pub fn derived_nullability(&self) -> Nullability { self.variants().any(|dtype| dtype.is_nullable()).into() } + + /// Returns a copy of these variants with `nullability` applied to every variant [`DType`]. + /// + /// A union has no top-level nullability of its own; its nullability is derived from its + /// variants (see [`Self::derived_nullability`]). Setting a union's nullability therefore means + /// propagating `nullability` into each variant. + pub fn with_nullability(&self, nullability: Nullability) -> Self { + let dtypes: Vec = self + .variants() + .map(|variant| variant.with_nullability(nullability)) + .collect(); + + Self::try_new(self.names().clone(), dtypes, self.type_ids().to_vec()) + .vortex_expect("changing variant nullability preserves union validity") + } } #[cfg(test)] @@ -526,16 +541,29 @@ mod tests { } #[test] - fn test_with_nullability_does_not_change_union_variants() { + fn test_with_nullability_propagates_to_variants() { let nonnullable = DType::Union(i32_variants()); - assert_eq!(nonnullable.as_nullable(), nonnullable); assert_eq!(nonnullable.nullability(), Nullability::NonNullable); - let nullable = DType::Union( - UnionVariants::new(["value"].into(), vec![DType::Utf8(Nullability::Nullable)]).unwrap(), - ); - assert_eq!(nullable.as_nonnullable(), nullable); + // `as_nullable` makes every variant nullable, since a union has no top-level nullability. + let nullable = nonnullable.as_nullable(); assert_eq!(nullable.nullability(), Nullability::Nullable); + assert_eq!( + nullable, + DType::Union( + UnionVariants::new( + ["int", "str"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + ], + ) + .unwrap() + ) + ); + + // `as_nonnullable` propagates the other direction, round-tripping back to the original. + assert_eq!(nullable.as_nonnullable(), nonnullable); } #[test]