Skip to content

Commit 357a4ba

Browse files
connortsui20connortsui20
authored
Remove stored Union Nullability (#8714)
## Summary Tracking issue: #7882 A union has no independent parent validity. A row is null when its selected child is null, so the union's logical nullability is determined by its variants. Storing the same information on `DType::Union` created two sources of truth and allowed the outer value to disagree with the variants. This change makes inconsistent union DTypes impossible. The remaining question is how an operation such as `mask` should change a union's variants when it introduces nulls. That remains deferred and does not block the canonical sparse `UnionArray`. The Arrow implementation notes and deferred operation scope are in [this comment on #7882](#7882 (comment)). ## Changes - change `DType::Union` to store only `UnionVariants` - derive union nullability at runtime from the variant DTypes - name the non-zero-cost scan `derived_nullability` - leave `with_nullability` unchanged for unions instead of rewriting variant schemas - restrict union least-supertype coercion to identical variants for now - remove Union nullability from Serde, protobuf, and FlatBuffers Supersedes #8713. --------- Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent d96e4ad commit 357a4ba

15 files changed

Lines changed: 176 additions & 237 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
289289
DType::Struct(fields, _) => fields
290290
.fields()
291291
.all(|field| supports_uncompressed_size_in_bytes(&field)),
292-
DType::Union(variants, _) => variants
292+
DType::Union(variants) => variants
293293
.variants()
294294
.all(|variant| supports_uncompressed_size_in_bytes(&variant)),
295295
DType::Variant(_) => false,

‎vortex-array/src/dtype/coercion.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,25 @@ impl DType {
7979
/// The core primitive — what type can hold both `self` and `other`?
8080
/// Returns `None` if no common supertype exists.
8181
pub fn least_supertype(&self, other: &DType) -> Option<DType> {
82+
match (self, other) {
83+
(DType::Union(lhs), DType::Union(rhs)) => {
84+
return (lhs == rhs).then(|| self.clone());
85+
}
86+
(DType::Null, DType::Union(variants)) => {
87+
return variants
88+
.derived_nullability()
89+
.is_nullable()
90+
.then(|| other.clone());
91+
}
92+
(DType::Union(variants), DType::Null) => {
93+
return variants
94+
.derived_nullability()
95+
.is_nullable()
96+
.then(|| self.clone());
97+
}
98+
_ => {}
99+
}
100+
82101
let union_null = self.nullability() | other.nullability();
83102

84103
if let (
@@ -308,6 +327,7 @@ mod tests {
308327

309328
use crate::dtype::DType;
310329
use crate::dtype::PType;
330+
use crate::dtype::UnionVariants;
311331
use crate::dtype::decimal::DecimalDType;
312332
use crate::dtype::nullability::Nullability::NonNullable;
313333
use crate::dtype::nullability::Nullability::Nullable;
@@ -349,6 +369,57 @@ mod tests {
349369
);
350370
}
351371

372+
#[test]
373+
fn least_supertype_union_requires_identical_variants() {
374+
let nonnullable = DType::Union(
375+
UnionVariants::new(
376+
["value"].into(),
377+
vec![DType::Primitive(PType::I32, NonNullable)],
378+
)
379+
.unwrap(),
380+
);
381+
let nullable = DType::Union(
382+
UnionVariants::new(
383+
["value"].into(),
384+
vec![DType::Primitive(PType::I32, Nullable)],
385+
)
386+
.unwrap(),
387+
);
388+
389+
assert_eq!(
390+
nonnullable.least_supertype(&nonnullable),
391+
Some(nonnullable.clone())
392+
);
393+
assert!(nonnullable.least_supertype(&nullable).is_none());
394+
assert!(nullable.least_supertype(&nonnullable).is_none());
395+
}
396+
397+
#[test]
398+
fn least_supertype_null_requires_nullable_union() {
399+
let nonnullable = DType::Union(
400+
UnionVariants::new(
401+
["value"].into(),
402+
vec![DType::Primitive(PType::I32, NonNullable)],
403+
)
404+
.unwrap(),
405+
);
406+
let nullable = DType::Union(
407+
UnionVariants::new(
408+
["value"].into(),
409+
vec![DType::Primitive(PType::I32, Nullable)],
410+
)
411+
.unwrap(),
412+
);
413+
414+
assert!(DType::Null.least_supertype(&nonnullable).is_none());
415+
assert!(nonnullable.least_supertype(&DType::Null).is_none());
416+
assert_eq!(
417+
DType::Null.least_supertype(&nullable),
418+
Some(nullable.clone())
419+
);
420+
assert_eq!(nullable.least_supertype(&DType::Null), Some(nullable));
421+
}
422+
352423
#[test]
353424
fn least_supertype_unsigned_widening() {
354425
let u8_nn = DType::Primitive(PType::U8, NonNullable);

‎vortex-array/src/dtype/dtype_impl.rs‎

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@ impl DType {
6464
| List(_, null)
6565
| FixedSizeList(_, _, null)
6666
| Struct(_, null)
67-
| Union(_, null)
6867
| Variant(null) => matches!(null, Nullability::Nullable),
68+
Union(variants) => variants.derived_nullability().is_nullable(),
6969
Extension(ext_dtype) => ext_dtype.storage_dtype().is_nullable(),
7070
}
7171
}
@@ -80,7 +80,10 @@ impl DType {
8080
self.with_nullability(Nullability::Nullable)
8181
}
8282

83-
/// Get a new DType with the given nullability (but otherwise the same as `self`)
83+
/// Get a new DType with the given nullability (but otherwise the same as `self`).
84+
///
85+
/// [`DType::Null`] and [`DType::Union`] have intrinsic nullability and are returned unchanged.
86+
/// To change a union's nullability, construct different [`UnionVariants`].
8487
pub fn with_nullability(&self, nullability: Nullability) -> Self {
8588
match self {
8689
Null => Null,
@@ -92,7 +95,7 @@ impl DType {
9295
List(edt, _) => List(Arc::clone(edt), nullability),
9396
FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability),
9497
Struct(sf, _) => Struct(sf.clone(), nullability),
95-
Union(vs, _) => Union(vs.clone(), nullability),
98+
Union(vs) => Union(vs.clone()),
9699
Variant(_) => Variant(nullability),
97100
Extension(ext) => Extension(ext.with_nullability(nullability)),
98101
}
@@ -124,7 +127,7 @@ impl DType {
124127
.zip_eq(rhs_dtype.fields())
125128
.all(|(l, r)| l.eq_ignore_nullability(&r)))
126129
}
127-
(Union(lhs, _), Union(rhs, _)) => {
130+
(Union(lhs), Union(rhs)) => {
128131
// Equal `names` implies equal length by FieldNames equality.
129132
lhs.names() == rhs.names()
130133
&& lhs.type_ids() == rhs.type_ids()
@@ -436,20 +439,12 @@ impl DType {
436439

437440
/// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise `None`.
438441
pub fn as_union_variants_opt(&self) -> Option<&UnionVariants> {
439-
if let Union(uv, _) = self {
440-
Some(uv)
441-
} else {
442-
None
443-
}
442+
if let Union(uv) = self { Some(uv) } else { None }
444443
}
445444

446445
/// Owned version of [Self::as_union_variants_opt].
447446
pub fn into_union_variants_opt(self) -> Option<UnionVariants> {
448-
if let Union(uv, _) = self {
449-
Some(uv)
450-
} else {
451-
None
452-
}
447+
if let Union(uv) = self { Some(uv) } else { None }
453448
}
454449

455450
/// Downcast a `DType` to an `ExtDType`
@@ -503,7 +498,7 @@ impl Display for DType {
503498
.map(|(field_null, dt)| format!("{field_null}={dt}"))
504499
.join(", "),
505500
),
506-
Union(uv, null) => write!(f, "union({uv}){null}"),
501+
Union(uv) => write!(f, "union({uv}){}", uv.derived_nullability()),
507502
Variant(null) => write!(f, "variant{null}"),
508503
Extension(ext) => write!(f, "{}", ext),
509504
}

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,13 @@ pub enum DType {
114114
/// A `Union` is composed of one or more **variants**, each with a name and a `DType`. A per-row
115115
/// `i8` tag selects which variant is "live" at that row.
116116
///
117+
/// Unlike other nested types, a union has no independent outer nullability. Its nullability is
118+
/// derived at runtime from its variants: a union is nullable when any variant can contain null.
119+
/// A concrete union array has no parent validity bitmap; a row's validity is the validity of
120+
/// its selected child.
121+
///
117122
/// See [`UnionVariants`] for the type-tag conventions and accessors.
118-
Union(UnionVariants, Nullability),
123+
Union(UnionVariants),
119124

120125
/// Dynamically typed values stored as Vortex scalars.
121126
///
@@ -148,7 +153,7 @@ impl PartialEq for DType {
148153
// StructFields handles its own Arc::ptr_eq in its PartialEq impl.
149154
(Self::Struct(a, na), Self::Struct(b, nb)) => na == nb && a == b,
150155
// UnionVariants handles its own Arc::ptr_eq in its PartialEq impl.
151-
(Self::Union(a, na), Self::Union(b, nb)) => na == nb && a == b,
156+
(Self::Union(a), Self::Union(b)) => a == b,
152157
(Self::Variant(a), Self::Variant(b)) => a == b,
153158
(Self::Extension(a), Self::Extension(b)) => a == b,
154159
// Every variant is listed in the first position so that adding a new

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

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ use itertools::Itertools;
1111
use vortex_error::VortexError;
1212
use vortex_error::VortexResult;
1313
use vortex_error::vortex_bail;
14-
use vortex_error::vortex_ensure;
1514
use vortex_error::vortex_err;
1615
use vortex_flatbuffers::FlatBuffer;
1716
use vortex_flatbuffers::FlatBufferRoot;
@@ -22,7 +21,6 @@ use vortex_session::VortexSession;
2221
use crate::dtype::DType;
2322
use crate::dtype::DecimalDType;
2423
use crate::dtype::FieldDType;
25-
use crate::dtype::Nullability;
2624
use crate::dtype::PType;
2725
use crate::dtype::StructFields;
2826
use crate::dtype::UnionVariants;
@@ -236,15 +234,7 @@ impl TryFrom<ViewedDType> for DType {
236234
.ok_or_else(|| vortex_err!("failed to parse union from flatbuffer"))?;
237235
let variants =
238236
UnionVariants::from_fb(fb_union, vfdt.buffer().clone(), vfdt.session.clone())?;
239-
240-
let nullability: Nullability = fb_union.nullable().into();
241-
vortex_ensure!(
242-
variants.nullability_constraints_satisfied(nullability),
243-
"Union nullability constraint not satisfied: nullability={:?}",
244-
nullability
245-
);
246-
247-
Ok(Self::Union(variants, nullability))
237+
Ok(Self::Union(variants))
248238
}
249239
fb::Type::Variant => {
250240
let fb_variant = fb
@@ -390,7 +380,7 @@ impl WriteFlatBuffer for DType {
390380
)
391381
.as_union_value()
392382
}
393-
Self::Union(uv, n) => {
383+
Self::Union(uv) => {
394384
let names = uv
395385
.names()
396386
.iter()
@@ -412,7 +402,6 @@ impl WriteFlatBuffer for DType {
412402
names,
413403
dtypes,
414404
type_ids,
415-
nullable: (*n).into(),
416405
},
417406
)
418407
.as_union_value()
@@ -583,7 +572,6 @@ mod test {
583572
],
584573
)
585574
.unwrap(),
586-
Nullability::NonNullable,
587575
)
588576
}
589577

@@ -593,14 +581,13 @@ mod test {
593581
}
594582

595583
#[test]
596-
fn test_union_round_trip_flatbuffer_with_nullability() {
584+
fn test_union_round_trip_flatbuffer_with_nullable_variant() {
597585
let dtype = DType::Union(
598586
UnionVariants::new(
599587
["null_variant", "str"].into(),
600588
vec![DType::Null, DType::Utf8(Nullability::NonNullable)],
601589
)
602590
.unwrap(),
603-
Nullability::Nullable,
604591
);
605592
roundtrip_dtype(dtype);
606593
}
@@ -618,7 +605,6 @@ mod test {
618605
vec![0, 5, 7],
619606
)
620607
.unwrap(),
621-
Nullability::NonNullable,
622608
);
623609

624610
let bytes = dtype.write_flatbuffer_bytes().unwrap();
@@ -631,7 +617,7 @@ mod test {
631617

632618
let deserialized = DType::try_from(view).unwrap();
633619
assert_eq!(dtype, deserialized);
634-
let DType::Union(uv, _) = &deserialized else {
620+
let DType::Union(uv) = &deserialized else {
635621
panic!("Expected Union");
636622
};
637623
assert_eq!(uv.type_ids(), &[0, 5, 7]);
@@ -654,7 +640,6 @@ mod test {
654640
vec![DType::Utf8(Nullability::NonNullable), struct_with_union],
655641
)
656642
.unwrap(),
657-
Nullability::NonNullable,
658643
);
659644

660645
roundtrip_dtype(outer_union);
@@ -710,7 +695,6 @@ mod test {
710695
names: Some(names),
711696
dtypes: Some(dtypes),
712697
type_ids: Some(type_ids),
713-
nullable: false,
714698
},
715699
);
716700

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ mod test {
2424
use crate::dtype::DType;
2525
use crate::dtype::Nullability;
2626
use crate::dtype::PType;
27+
use crate::dtype::UnionVariants;
2728
use crate::dtype::serde::DTypeSerde;
2829
use crate::dtype::test::SESSION;
2930

@@ -181,4 +182,24 @@ mod test {
181182
.unwrap();
182183
assert_eq!(DType::Variant(Nullability::Nullable), deserialized);
183184
}
185+
186+
#[test]
187+
fn test_serde_union_dtype_json_roundtrip() {
188+
let dtype = DType::Union(
189+
UnionVariants::new(["value"].into(), vec![DType::Utf8(Nullability::Nullable)]).unwrap(),
190+
);
191+
192+
let json = serde_json::to_string(&dtype).unwrap();
193+
assert_eq!(
194+
json,
195+
r#"{"Union":{"names":["value"],"dtypes":[{"Utf8":true}],"type_ids":[0]}}"#
196+
);
197+
let mut deserializer = serde_json::Deserializer::from_str(&json);
198+
let deserialized: DType = DTypeSerde::<DType>::new(&SESSION)
199+
.deserialize(&mut deserializer)
200+
.unwrap();
201+
202+
assert_eq!(deserialized, dtype);
203+
assert_eq!(deserialized.nullability(), Nullability::Nullable);
204+
}
184205
}

0 commit comments

Comments
 (0)