Add outer nullability to Union dtypes

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs

@@ -289,7 +289,7 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
         DType::Struct(fields, _) => fields
             .fields()
             .all(|field| supports_uncompressed_size_in_bytes(&field)),
-        DType::Union(variants) => variants
+        DType::Union(variants, _) => variants
             .variants()
             .all(|variant| supports_uncompressed_size_in_bytes(&variant)),
         DType::Variant(_) => false,

vortex-array/src/dtype/coercion.rs

@@ -80,20 +80,14 @@ impl DType {
     /// Returns `None` if no common supertype exists.
     pub fn least_supertype(&self, other: &DType) -> Option<DType> {
         match (self, other) {
-            (DType::Union(lhs), DType::Union(rhs)) => {
-                return (lhs == rhs).then(|| self.clone());
+            (DType::Union(lhs, lhs_null), DType::Union(rhs, rhs_null)) => {
+                return (lhs == rhs).then(|| DType::Union(lhs.clone(), *lhs_null | *rhs_null));
             }
-            (DType::Null, DType::Union(variants)) => {
-                return variants
-                    .derived_nullability()
-                    .is_nullable()
-                    .then(|| other.clone());
+            (DType::Null, DType::Union(..)) => {
+                return Some(other.as_nullable());
             }
-            (DType::Union(variants), DType::Null) => {
-                return variants
-                    .derived_nullability()
-                    .is_nullable()
-                    .then(|| self.clone());
+            (DType::Union(..), DType::Null) => {
+                return Some(self.as_nullable());
             }
             _ => {}
         }
@@ -377,13 +371,15 @@ mod tests {
                 vec![DType::Primitive(PType::I32, NonNullable)],
             )
             .unwrap(),
+            NonNullable,
         );
         let nullable = DType::Union(
             UnionVariants::new(
                 ["value"].into(),
                 vec![DType::Primitive(PType::I32, Nullable)],
             )
             .unwrap(),
+            NonNullable,
         );
 
         assert_eq!(
@@ -395,29 +391,22 @@ mod tests {
     }
 
     #[test]
-    fn least_supertype_null_requires_nullable_union() {
+    fn least_supertype_null_makes_union_outer_nullable() {
         let nonnullable = DType::Union(
             UnionVariants::new(
                 ["value"].into(),
                 vec![DType::Primitive(PType::I32, NonNullable)],
             )
             .unwrap(),
+            NonNullable,
         );
-        let nullable = DType::Union(
-            UnionVariants::new(
-                ["value"].into(),
-                vec![DType::Primitive(PType::I32, Nullable)],
-            )
-            .unwrap(),
-        );
+        let nullable = nonnullable.as_nullable();
 
-        assert!(DType::Null.least_supertype(&nonnullable).is_none());
-        assert!(nonnullable.least_supertype(&DType::Null).is_none());
         assert_eq!(
-            DType::Null.least_supertype(&nullable),
+            DType::Null.least_supertype(&nonnullable),
             Some(nullable.clone())
         );
-        assert_eq!(nullable.least_supertype(&DType::Null), Some(nullable));
+        assert_eq!(nonnullable.least_supertype(&DType::Null), Some(nullable));
     }
 
     #[test]

vortex-array/src/dtype/dtype_impl.rs

@@ -64,8 +64,8 @@ impl DType {
             | List(_, null)
             | FixedSizeList(_, _, null)
             | Struct(_, null)
+            | Union(_, null)
             | Variant(null) => matches!(null, Nullability::Nullable),
-            Union(variants) => variants.derived_nullability().is_nullable(),
             Extension(ext_dtype) => ext_dtype.storage_dtype().is_nullable(),
         }
     }
@@ -82,8 +82,7 @@ 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.
     pub fn with_nullability(&self, nullability: Nullability) -> Self {
         match self {
             Null => Null,
@@ -95,7 +94,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.clone(), nullability),
             Variant(_) => Variant(nullability),
             Extension(ext) => Extension(ext.with_nullability(nullability)),
         }
@@ -127,7 +126,7 @@ impl DType {
                         .zip_eq(rhs_dtype.fields())
                         .all(|(l, r)| l.eq_ignore_nullability(&r)))
             }
-            (Union(lhs), Union(rhs)) => {
+            (Union(lhs, _), Union(rhs, _)) => {
                 // Equal `names` implies equal length by FieldNames equality.
                 lhs.names() == rhs.names()
                     && lhs.type_ids() == rhs.type_ids()
@@ -439,12 +438,20 @@ impl DType {
 
     /// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise `None`.
     pub fn as_union_variants_opt(&self) -> Option<&UnionVariants> {
-        if let Union(uv) = self { Some(uv) } else { None }
+        if let Union(uv, _) = self {
+            Some(uv)
+        } else {
+            None
+        }
     }
 
     /// Owned version of [Self::as_union_variants_opt].
     pub fn into_union_variants_opt(self) -> Option<UnionVariants> {
-        if let Union(uv) = self { Some(uv) } else { None }
+        if let Union(uv, _) = self {
+            Some(uv)
+        } else {
+            None
+        }
     }
 
     /// Downcast a `DType` to an `ExtDType`
@@ -498,7 +505,7 @@ impl Display for DType {
                     .map(|(field_null, dt)| format!("{field_null}={dt}"))
                     .join(", "),
             ),
-            Union(uv) => write!(f, "union({uv}){}", uv.derived_nullability()),
+            Union(uv, null) => write!(f, "union({uv}){null}"),
             Variant(null) => write!(f, "variant{null}"),
             Extension(ext) => write!(f, "{}", ext),
         }

vortex-array/src/dtype/mod.rs

@@ -111,15 +111,13 @@ pub enum DType {
     /// A logical union (sum) type.
     ///
     /// A `Union` is composed of one or more **variants**, each with a name and a `DType`. A per-row
-    /// `i8` tag selects which variant is "live" at that row.
+    /// `u8` tag selects which variant is "live" at that row.
     ///
-    /// Unlike other nested types, a union has no independent outer nullability. Its nullability is
-    /// derived at runtime from its variants: a union is nullable when any variant can contain null.
-    /// A concrete union array has no parent validity bitmap; a row's validity is the validity of
-    /// its selected child.
+    /// A union has independent outer nullability in addition to the nullability of each variant.
+    /// This distinguishes a null union from a non-null union whose selected child is null.
     ///
     /// See [`UnionVariants`] for the type-tag conventions and accessors.
-    Union(UnionVariants),
+    Union(UnionVariants, Nullability),
 
     /// Dynamically typed values stored as Vortex scalars.
     ///
@@ -152,7 +150,7 @@ impl PartialEq for DType {
             // StructFields handles its own Arc::ptr_eq in its PartialEq impl.
             (Self::Struct(a, na), Self::Struct(b, nb)) => na == nb && a == b,
             // UnionVariants handles its own Arc::ptr_eq in its PartialEq impl.
-            (Self::Union(a), Self::Union(b)) => a == b,
+            (Self::Union(a, na), Self::Union(b, nb)) => na == nb && a == b,
             (Self::Variant(a), Self::Variant(b)) => a == b,
             (Self::Extension(a), Self::Extension(b)) => a == b,
             // Every variant is listed in the first position so that adding a new

vortex-array/src/dtype/serde/flatbuffers.rs

@@ -235,7 +235,7 @@ impl TryFrom<ViewedDType> for DType {
                     .ok_or_else(|| vortex_err!("failed to parse union from flatbuffer"))?;
                 let variants =
                     UnionVariants::from_fb(fb_union, vfdt.buffer().clone(), vfdt.session.clone())?;
-                Ok(Self::Union(variants))
+                Ok(Self::Union(variants, fb_union.nullable().into()))
             }
             fb::Type::Variant => {
                 let fb_variant = fb
@@ -381,7 +381,7 @@ impl WriteFlatBuffer for DType {
                 )
                 .as_union_value()
             }
-            Self::Union(uv) => {
+            Self::Union(uv, n) => {
                 let names = uv
                     .names()
                     .iter()
@@ -411,6 +411,7 @@ impl WriteFlatBuffer for DType {
                         names,
                         dtypes,
                         type_ids,
+                        nullable: (*n).into(),
                     },
                 )
                 .as_union_value()
@@ -581,6 +582,7 @@ mod test {
                 ],
             )
             .unwrap(),
+            Nullability::NonNullable,
         )
     }
 
@@ -597,6 +599,7 @@ mod test {
                 vec![DType::Null, DType::Utf8(Nullability::NonNullable)],
             )
             .unwrap(),
+            Nullability::NonNullable,
         );
         roundtrip_dtype(dtype);
     }
@@ -614,6 +617,7 @@ mod test {
                 vec![0, 5, u8::MAX],
             )
             .unwrap(),
+            Nullability::Nullable,
         );
 
         let bytes = dtype.write_flatbuffer_bytes().unwrap();
@@ -626,7 +630,7 @@ mod test {
 
         let deserialized = DType::try_from(view).unwrap();
         assert_eq!(dtype, deserialized);
-        let DType::Union(uv) = &deserialized else {
+        let DType::Union(uv, _) = &deserialized else {
             panic!("Expected Union");
         };
         assert_eq!(uv.type_ids(), &[0, 5, u8::MAX]);
@@ -649,6 +653,7 @@ mod test {
                 vec![DType::Utf8(Nullability::NonNullable), struct_with_union],
             )
             .unwrap(),
+            Nullability::Nullable,
         );
 
         roundtrip_dtype(outer_union);
@@ -704,6 +709,7 @@ mod test {
                 names: Some(names),
                 dtypes: Some(dtypes),
                 type_ids: Some(type_ids),
+                nullable: false,
             },
         );
 

vortex-array/src/dtype/serde/mod.rs

@@ -187,12 +187,13 @@ mod test {
     fn test_serde_union_dtype_json_roundtrip() {
         let dtype = DType::Union(
             UnionVariants::new(["value"].into(), vec![DType::Utf8(Nullability::Nullable)]).unwrap(),
+            Nullability::Nullable,
         );
 
         let json = serde_json::to_string(&dtype).unwrap();
         assert_eq!(
             json,
-            r#"{"Union":{"names":["value"],"dtypes":[{"Utf8":true}],"type_ids":[0]}}"#
+            r#"{"Union":[{"names":["value"],"dtypes":[{"Utf8":true}],"type_ids":[0]},true]}"#
         );
         let mut deserializer = serde_json::Deserializer::from_str(&json);
         let deserialized: DType = DTypeSerde::<DType>::new(&SESSION)

vortex-array/src/dtype/serde/proto.rs

@@ -104,7 +104,7 @@ impl DType {
                     })
                     .collect::<VortexResult<Vec<_>>>()?;
                 let variants = UnionVariants::try_new(names, dtypes, type_ids)?;
-                Ok(Self::Union(variants))
+                Ok(Self::Union(variants, u.nullable.into()))
             }
             DtypeType::Variant(v) => Ok(Self::Variant(v.nullable.into())),
             DtypeType::Extension(e) => {
@@ -173,13 +173,14 @@ impl TryFrom<&DType> for pb::DType {
                         .collect::<VortexResult<Vec<_>>>()?,
                     nullable: (*null).into(),
                 }),
-                DType::Union(uv) => DtypeType::Union(pb::Union {
+                DType::Union(uv, null) => DtypeType::Union(pb::Union {
                     names: uv.names().iter().map(|n| n.as_ref().to_string()).collect(),
                     dtypes: uv
                         .variants()
                         .map(|d| Self::try_from(&d))
                         .collect::<VortexResult<Vec<_>>>()?,
                     type_ids: uv.type_ids().iter().map(|t| *t as i32).collect(),
+                    nullable: (*null).into(),
                 }),
                 DType::Variant(null) => DtypeType::Variant(pb::Variant {
                     nullable: (*null).into(),
@@ -416,22 +417,20 @@ mod tests {
             ],
         )
         .unwrap();
-        let dtype = DType::Union(variants);
+        let dtype = DType::Union(variants, Nullability::NonNullable);
         let converted = round_trip_dtype(&dtype);
         assert_eq!(dtype, converted);
     }
 
     #[test]
-    fn test_union_round_trip_proto_with_nullable_variant() {
+    fn test_union_round_trip_proto_with_inner_and_outer_nullability() {
         let variants = UnionVariants::new(
             ["null_variant", "str"].into(),
             vec![DType::Null, DType::Utf8(Nullability::NonNullable)],
         )
         .unwrap();
 
-        assert_eq!(variants.derived_nullability(), Nullability::Nullable);
-
-        let dtype = DType::Union(variants);
+        let dtype = DType::Union(variants, Nullability::Nullable);
         let converted = round_trip_dtype(&dtype);
         assert_eq!(dtype, converted);
     }
@@ -449,11 +448,11 @@ mod tests {
         )
         .unwrap();
 
-        let dtype = DType::Union(variants);
+        let dtype = DType::Union(variants, Nullability::Nullable);
         let converted = round_trip_dtype(&dtype);
         assert_eq!(dtype, converted);
 
-        let DType::Union(uv) = &converted else {
+        let DType::Union(uv, _) = &converted else {
             panic!("Expected Union");
         };
         assert_eq!(uv.type_ids(), &[0, 5, 7]);
@@ -471,6 +470,7 @@ mod tests {
                 ],
                 // 256 does not fit in u8.
                 type_ids: vec![256],
+                nullable: false,
             })),
         };
 
@@ -495,6 +495,7 @@ mod tests {
                 ],
             )
             .unwrap(),
+            Nullability::NonNullable,
         );
 
         let struct_with_union = DType::Struct(
@@ -511,6 +512,7 @@ mod tests {
                 vec![DType::Utf8(Nullability::NonNullable), struct_with_union],
             )
             .unwrap(),
+            Nullability::Nullable,
         );
 
         let converted = round_trip_dtype(&outer_union);