From 2eb8864027e84c3ad3345a85bed4c32090e73f11 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 16 Jul 2026 10:25:00 -0400 Subject: [PATCH] Use u8 Union type IDs Signed-off-by: Connor Tsui --- vortex-array/src/dtype/serde/flatbuffers.rs | 17 +++- vortex-array/src/dtype/serde/proto.rs | 10 +-- vortex-array/src/dtype/serde/serde.rs | 2 +- vortex-array/src/dtype/union.rs | 89 +++++++++---------- .../flatbuffers/vortex-dtype/dtype.fbs | 2 +- vortex-proto/proto/dtype.proto | 2 +- vortex-proto/src/generated/vortex.dtype.rs | 2 +- 7 files changed, 65 insertions(+), 59 deletions(-) diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index eef15d59c57..6610ad5a580 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -117,10 +117,11 @@ impl UnionVariants { }) .collect::>(); - let type_ids: Vec = fb_union + let type_ids: Vec = fb_union .type_ids() .ok_or_else(|| vortex_err!("failed to parse union type_ids from flatbuffer"))? .iter() + .map(i8::cast_unsigned) .collect(); UnionVariants::try_from_fields(names, dtypes, type_ids) @@ -394,7 +395,15 @@ impl WriteFlatBuffer for DType { .collect::>>()?; let dtypes = Some(fbb.create_vector(&dtypes)); - let type_ids = Some(fbb.create_vector(uv.type_ids())); + // The FlatBuffers schema retains its original signed byte wire type for backwards + // compatibility. Reinterpreting the bits preserves the full u8 tag range. + let type_ids = uv + .type_ids() + .iter() + .copied() + .map(u8::cast_signed) + .collect_vec(); + let type_ids = Some(fbb.create_vector(&type_ids)); fb::Union::create( fbb, @@ -602,7 +611,7 @@ mod test { DType::Utf8(Nullability::NonNullable), DType::Bool(Nullability::NonNullable), ], - vec![0, 5, 7], + vec![0, 5, u8::MAX], ) .unwrap(), ); @@ -620,7 +629,7 @@ mod test { let DType::Union(uv) = &deserialized else { panic!("Expected Union"); }; - assert_eq!(uv.type_ids(), &[0, 5, 7]); + assert_eq!(uv.type_ids(), &[0, 5, u8::MAX]); } #[test] diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index e3dfb90557a..a0c27cb298b 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -98,8 +98,8 @@ impl DType { .type_ids .iter() .map(|t| { - i8::try_from(*t).map_err(|_| { - vortex_err!("Union type_id {t} somehow does not fit in i8") + u8::try_from(*t).map_err(|_| { + vortex_err!("Union type_id {t} somehow does not fit in u8") }) }) .collect::>>()?; @@ -469,8 +469,8 @@ mod tests { pb::DType::try_from(&DType::Primitive(PType::I32, Nullability::NonNullable)) .unwrap(), ], - // 200 does not fit in i8. - type_ids: vec![200], + // 256 does not fit in u8. + type_ids: vec![256], })), }; @@ -480,7 +480,7 @@ mod tests { result .unwrap_err() .to_string() - .contains("does not fit in i8") + .contains("does not fit in u8") ); } diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index 1d386a34cbf..7f360382d7e 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -431,7 +431,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, UnionVariants> { { let mut names: Option = None; let mut dtypes: Option> = None; - let mut type_ids: Option> = None; + let mut type_ids: Option> = None; while let Some(key) = map.next_key::>()? { match key.as_ref() { diff --git a/vortex-array/src/dtype/union.rs b/vortex-array/src/dtype/union.rs index faec75751c5..b14a3b84d21 100644 --- a/vortex-array/src/dtype/union.rs +++ b/vortex-array/src/dtype/union.rs @@ -19,13 +19,11 @@ use crate::dtype::Nullability; /// Type information for a union array. /// /// A `UnionVariants` describes the possible alternative types for each row of a union, along with a -/// per-variant `i8` type tag. We use the term **variants** (rather than "fields") because a union +/// per-variant `u8` type tag. We use the term **variants** (rather than "fields") because a union /// is a sum type: each row chooses exactly one alternative. /// -/// Per Arrow's spec, the per-row type tag is an `int8`. By default, tag `i` selects the child at -/// offset `i` (`type_ids = [0, 1, ..., N-1]`). -/// -/// Type IDs are restricted to the Arrow-compatible range `0..=127`; negative tags are rejected. +/// By default, tag `i` selects the child at offset `i` (`type_ids = [0, 1, ..., N-1]`). Vortex uses +/// unsigned tags and supports up to 256 variants. /// /// Schemas may also use non-consecutive tags (e.g. `[0, 5, 7]`), in which case the value of /// `type_ids[i]` is the tag used in the data to select the child at offset `i`. Supporting @@ -114,11 +112,11 @@ struct UnionVariantsInner { /// consecutive offsets is `[0, 1, ..., N-1]`. /// /// For schemas with explicit `typeIds` indirection (e.g. `[0, 5, 7]`), this stores those tags. - type_ids: Arc<[i8]>, + type_ids: Arc<[u8]>, } impl UnionVariantsInner { - fn from_fields(names: FieldNames, dtypes: Arc<[FieldDType]>, type_ids: Arc<[i8]>) -> Self { + fn from_fields(names: FieldNames, dtypes: Arc<[FieldDType]>, type_ids: Arc<[u8]>) -> Self { Self { names, dtypes, @@ -160,7 +158,7 @@ impl UnionVariants { } /// Validate that `names`, `dtypes`, and `type_ids` are mutually consistent. - fn validate_shape(names: &FieldNames, n_dtypes: usize, type_ids: &[i8]) -> VortexResult<()> { + fn validate_shape(names: &FieldNames, n_dtypes: usize, type_ids: &[u8]) -> VortexResult<()> { vortex_ensure_eq!( names.len(), n_dtypes, @@ -182,9 +180,10 @@ impl UnionVariants { type_ids ); vortex_ensure!( - type_ids.iter().all(|type_id| *type_id >= 0), - "type_ids must be non-negative, got {:?}", - type_ids + type_ids.len() <= u8::MAX as usize + 1, + "union supports at most {} variants, got {}", + u8::MAX as usize + 1, + type_ids.len() ); vortex_ensure!( names.iter().all_unique(), @@ -200,12 +199,12 @@ impl UnionVariants { /// # Errors /// /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there - /// are any duplicate names or type ids, or if a type ID is negative. - pub fn try_new(names: FieldNames, dtypes: Vec, type_ids: Vec) -> VortexResult { + /// are any duplicate names or type IDs, or if there are more than 256 variants. + pub fn try_new(names: FieldNames, dtypes: Vec, type_ids: Vec) -> VortexResult { Self::validate_shape(&names, dtypes.len(), &type_ids)?; let dtypes: Arc<[FieldDType]> = dtypes.into_iter().map(FieldDType::from).collect(); - let type_ids: Arc<[i8]> = Arc::from(type_ids); + let type_ids: Arc<[u8]> = Arc::from(type_ids); Ok(Self(Arc::new(UnionVariantsInner::from_fields( names, dtypes, type_ids, @@ -217,21 +216,21 @@ impl UnionVariants { /// # Errors /// /// `names` and `dtypes` must have the same length, and `names.len()` cannot be more than - /// `i8::MAX as usize + 1` (128). + /// `u8::MAX as usize + 1` (256). pub fn new(names: FieldNames, dtypes: Vec) -> VortexResult { - const MAX_CONSECUTIVE: usize = i8::MAX as usize + 1; + const MAX_VARIANTS: usize = u8::MAX as usize + 1; vortex_ensure!( - names.len() <= MAX_CONSECUTIVE, + names.len() <= MAX_VARIANTS, "union supports at most {} consecutive variants, got {}", - MAX_CONSECUTIVE, + MAX_VARIANTS, names.len() ); #[expect( clippy::cast_possible_truncation, - reason = "the MAX_CONSECUTIVE bound above guarantees `i as i8` is in range" + reason = "the MAX_VARIANTS bound above guarantees `i as u8` is in range" )] - let type_ids: Vec = (0..names.len()).map(|i| i as i8).collect(); + let type_ids: Vec = (0..names.len()).map(|i| i as u8).collect(); Self::try_new(names, dtypes, type_ids) } @@ -244,11 +243,11 @@ impl UnionVariants { /// # Errors /// /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there - /// are any duplicate names or type ids, or if a type ID is negative. + /// are any duplicate names or type IDs, or if there are more than 256 variants. pub(crate) fn try_from_fields( names: FieldNames, dtypes: Vec, - type_ids: Vec, + type_ids: Vec, ) -> VortexResult { Self::validate_shape(&names, dtypes.len(), &type_ids)?; @@ -276,7 +275,7 @@ impl UnionVariants { /// Returns the per-variant type tag vector. Entry `i` is the tag that the data uses to /// select the variant at offset `i`. - pub fn type_ids(&self) -> &[i8] { + pub fn type_ids(&self) -> &[u8] { &self.0.type_ids } @@ -286,12 +285,12 @@ impl UnionVariants { .type_ids .iter() .enumerate() - .all(|(i, &tag)| i8::try_from(i).is_ok_and(|i| i == tag)) + .all(|(i, &tag)| u8::try_from(i).is_ok_and(|i| i == tag)) } /// Find the offset of a variant by name. Returns `None` if no variant has the name. /// - /// This is a linear scan over [`Self::names`]. A union has at most 128 variants (and usually + /// This is a linear scan over [`Self::names`]. A union has at most 256 variants (and usually /// far fewer), so a linear scan beats building and probing a `HashMap` in practice. pub fn find(&self, name: impl AsRef) -> Option { let name = name.as_ref(); @@ -331,9 +330,9 @@ impl UnionVariants { /// `type_ids`. /// /// This is a linear scan over [`Self::type_ids`]. The number of variants is bounded by - /// `i8::MAX + 1 = 128`, which fits in two cache lines, so a linear scan is faster than a - /// `HashMap` lookup in practice. - pub fn tag_to_child_index(&self, tag: i8) -> Option { + /// 256, so a linear scan is faster than a `HashMap` lookup in practice for the + /// small unions typically encountered. + pub fn tag_to_child_index(&self, tag: u8) -> Option { self.0.type_ids.iter().position(|&t| t == tag) } @@ -342,7 +341,7 @@ impl UnionVariants { /// # Panics /// /// Panics if `index >= self.len()`. - pub fn child_index_to_tag(&self, index: usize) -> i8 { + pub fn child_index_to_tag(&self, index: usize) -> u8 { self.0.type_ids[index] } @@ -437,14 +436,14 @@ mod tests { } #[test] - fn test_negative_type_ids_rejected() { - let result = UnionVariants::try_new( - ["negative"].into(), + fn test_high_type_id() { + let variants = UnionVariants::try_new( + ["high"].into(), vec![DType::Primitive(PType::I32, Nullability::NonNullable)], - vec![-1], - ); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("non-negative")); + vec![u8::MAX], + ) + .unwrap(); + assert_eq!(variants.type_ids(), &[u8::MAX]); } #[test] @@ -575,23 +574,21 @@ mod tests { #[test] fn test_new_max_size() { - // 128 variants is the maximum for consecutive type_ids: tags 0..=127 all fit in i8. - let names: Vec = (0..128).map(|i| format!("v{i}")).collect(); - let dtypes: Vec = (0..128) + let names: Vec = (0..256).map(|i| format!("v{i}")).collect(); + let dtypes: Vec = (0..256) .map(|_| DType::Primitive(PType::I32, Nullability::NonNullable)) .collect(); let names: FieldNames = names.into_iter().collect(); let variants = UnionVariants::new(names, dtypes).unwrap(); - assert_eq!(variants.len(), 128); - assert_eq!(variants.type_ids()[127], 127); + assert_eq!(variants.len(), 256); + assert_eq!(variants.type_ids()[255], 255); assert!(variants.is_consecutive()); } #[test] fn test_new_too_large_rejected() { - // 129 variants exceeds i8::MAX + 1 = 128. - let names: Vec = (0..129).map(|i| format!("v{i}")).collect(); - let dtypes: Vec = (0..129) + let names: Vec = (0..257).map(|i| format!("v{i}")).collect(); + let dtypes: Vec = (0..257) .map(|_| DType::Primitive(PType::I32, Nullability::NonNullable)) .collect(); let names: FieldNames = names.into_iter().collect(); @@ -601,7 +598,7 @@ mod tests { result .unwrap_err() .to_string() - .contains("at most 128 consecutive variants") + .contains("at most 256 consecutive variants") ); } @@ -610,7 +607,7 @@ mod tests { let v = UnionVariants::empty(); assert!(v.is_empty()); assert_eq!(v.len(), 0); - assert_eq!(v.type_ids(), &[] as &[i8]); + assert_eq!(v.type_ids(), &[] as &[u8]); assert!(v.is_consecutive()); } } diff --git a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs index add1ae63e2e..f3b51c845b1 100644 --- a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs +++ b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs @@ -70,7 +70,7 @@ table Variant { table Union { names: [string]; dtypes: [DType]; - type_ids: [byte]; // length must equal dtypes.len() + type_ids: [byte]; // length must equal dtypes.len(); interpreted as unsigned } union Type { diff --git a/vortex-proto/proto/dtype.proto b/vortex-proto/proto/dtype.proto index decc2179080..675a9156f22 100644 --- a/vortex-proto/proto/dtype.proto +++ b/vortex-proto/proto/dtype.proto @@ -77,7 +77,7 @@ message Variant { message Union { repeated string names = 1; repeated DType dtypes = 2; - repeated int32 type_ids = 3; // length must equal dtypes.len(); each value must fit in int8 + repeated int32 type_ids = 3; // length must equal dtypes.len(); each value must fit in uint8 } message DType { diff --git a/vortex-proto/src/generated/vortex.dtype.rs b/vortex-proto/src/generated/vortex.dtype.rs index 4bffca278e2..984b4b6b71e 100644 --- a/vortex-proto/src/generated/vortex.dtype.rs +++ b/vortex-proto/src/generated/vortex.dtype.rs @@ -77,7 +77,7 @@ pub struct Union { pub names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, #[prost(message, repeated, tag = "2")] pub dtypes: ::prost::alloc::vec::Vec, - /// length must equal dtypes.len(); each value must fit in int8 + /// length must equal dtypes.len(); each value must fit in uint8 #[prost(int32, repeated, tag = "3")] pub type_ids: ::prost::alloc::vec::Vec, }