Make union as_nullable propagate to all variants

A union has no top-level nullability of its own; its nullability is
derived from its variants. So DType::with_nullability on a union now
propagates the requested nullability into every variant instead of being
a no-op, via a new UnionVariants::with_nullability helper.

This also lets Scalar::into_nullable work on non-nullable unions (it
previously had to assert against the no-nullable-representation case), so
the assertion and its panic path are removed.

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: connortsui20

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)),
         }

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<DType> = 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]

vortex-array/src/scalar/cast.rs

@@ -76,19 +76,11 @@ impl Scalar {
 
     /// Cast the scalar into a nullable version of its current type.
     ///
-    /// # Panics
-    ///
-    /// Panics if the scalar's dtype has no nullable representation, such as a union dtype when all
-    /// of its variants are non-nullable.
+    /// 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();
-        let nullable_dtype = dtype.as_nullable();
-        assert!(
-            nullable_dtype.is_nullable(),
-            "dtype {dtype} cannot be made nullable"
-        );
-
-        Self::try_new(nullable_dtype, value)
+        Self::try_new(dtype.as_nullable(), value)
             .vortex_expect("Casting to nullable should always succeed")
     }
 }

vortex-array/src/scalar/tests/casting.rs

@@ -416,50 +416,39 @@ mod tests {
     }
 
     #[test]
-    fn into_nullable_runtime_nullable_union_is_identity() -> VortexResult<()> {
-        let variants = UnionVariants::try_new(
-            ["int", "string"].into(),
-            vec![
-                DType::Primitive(PType::I32, Nullability::Nullable),
-                DType::Utf8(Nullability::NonNullable),
-            ],
-            vec![5, 9],
-        )?;
+    fn into_nullable_makes_every_variant_nullable() -> VortexResult<()> {
         let scalar = Scalar::union(
-            variants,
+            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::Nullable),
+            Scalar::primitive(42_i32, Nullability::NonNullable),
         )?;
-        let dtype = scalar.dtype().clone();
 
         let nullable = scalar.into_nullable();
 
-        assert_eq!(nullable.dtype(), &dtype);
-        assert!(nullable.dtype().is_nullable());
-        assert_eq!(nullable.as_union().type_id(), Some(5));
-
-        Ok(())
-    }
-
-    #[test]
-    #[should_panic(expected = "cannot be made nullable")]
-    fn into_nullable_runtime_nonnullable_union_panics() {
-        let variants = UnionVariants::try_new(
+        // 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::NonNullable),
-                DType::Utf8(Nullability::NonNullable),
+                DType::Primitive(PType::I32, Nullability::Nullable),
+                DType::Utf8(Nullability::Nullable),
             ],
             vec![5, 9],
-        )
-        .vortex_expect("valid union variants");
-        let scalar = Scalar::union(
-            variants,
-            5,
-            Scalar::primitive(42_i32, Nullability::NonNullable),
-        )
-        .vortex_expect("valid union scalar");
+        )?);
+        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))
+        );
 
-        scalar.into_nullable();
+        Ok(())
     }
 }