Use u8 Union type IDs

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

Author: connortsui20

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

@@ -117,7 +117,7 @@ impl UnionVariants {
             })
             .collect::<Vec<_>>();
 
-        let type_ids: Vec<i8> = fb_union
+        let type_ids: Vec<u8> = fb_union
             .type_ids()
             .ok_or_else(|| vortex_err!("failed to parse union type_ids from flatbuffer"))?
             .iter()
@@ -687,7 +687,7 @@ mod test {
             .unwrap();
         let dtypes = fbb.create_vector(&[dtype_a, dtype_b]);
 
-        let type_ids = fbb.create_vector::<i8>(&[0]);
+        let type_ids = fbb.create_vector::<u8>(&[0]);
 
         let union_table = fb::Union::create(
             &mut fbb,

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::<VortexResult<Vec<_>>>()?;
@@ -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")
         );
     }
 

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

@@ -431,7 +431,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, UnionVariants> {
             {
                 let mut names: Option<FieldNames> = None;
                 let mut dtypes: Option<Vec<DType>> = None;
-                let mut type_ids: Option<Vec<i8>> = None;
+                let mut type_ids: Option<Vec<u8>> = None;
 
                 while let Some(key) = map.next_key::<Cow<'_, str>>()? {
                     match key.as_ref() {

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 255 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,
+            "union supports at most {} variants, got {}",
+            u8::MAX,
+            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<DType>, type_ids: Vec<i8>) -> VortexResult<Self> {
+    /// are any duplicate names or type IDs, or if there are more than 255 variants.
+    pub fn try_new(names: FieldNames, dtypes: Vec<DType>, type_ids: Vec<u8>) -> VortexResult<Self> {
         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` (255).
     pub fn new(names: FieldNames, dtypes: Vec<DType>) -> VortexResult<Self> {
-        const MAX_CONSECUTIVE: usize = i8::MAX as usize + 1;
+        const MAX_VARIANTS: usize = u8::MAX as usize;
         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<i8> = (0..names.len()).map(|i| i as i8).collect();
+        let type_ids: Vec<u8> = (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 255 variants.
     pub(crate) fn try_from_fields(
         names: FieldNames,
         dtypes: Vec<FieldDType>,
-        type_ids: Vec<i8>,
+        type_ids: Vec<u8>,
     ) -> VortexResult<Self> {
         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 255 variants (and usually
     /// far fewer), so a linear scan beats building and probing a `HashMap` in practice.
     pub fn find(&self, name: impl AsRef<str>) -> Option<usize> {
         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<usize> {
+    /// `u8::MAX = 255`, 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<usize> {
         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,22 @@ 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<String> = (0..128).map(|i| format!("v{i}")).collect();
-        let dtypes: Vec<DType> = (0..128)
+        // 255 variants is the maximum supported size.
+        let names: Vec<String> = (0..255).map(|i| format!("v{i}")).collect();
+        let dtypes: Vec<DType> = (0..255)
             .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(), 255);
+        assert_eq!(variants.type_ids()[254], 254);
         assert!(variants.is_consecutive());
     }
 
     #[test]
     fn test_new_too_large_rejected() {
-        // 129 variants exceeds i8::MAX + 1 = 128.
-        let names: Vec<String> = (0..129).map(|i| format!("v{i}")).collect();
-        let dtypes: Vec<DType> = (0..129)
+        let names: Vec<String> = (0..256).map(|i| format!("v{i}")).collect();
+        let dtypes: Vec<DType> = (0..256)
             .map(|_| DType::Primitive(PType::I32, Nullability::NonNullable))
             .collect();
         let names: FieldNames = names.into_iter().collect();
@@ -601,7 +599,7 @@ mod tests {
             result
                 .unwrap_err()
                 .to_string()
-                .contains("at most 128 consecutive variants")
+                .contains("at most 255 consecutive variants")
         );
     }
 
@@ -610,7 +608,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());
     }
 }

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: [ubyte]; // length must equal dtypes.len()
 }
 
 union Type {

vortex-flatbuffers/src/generated/dtype.rs

@@ -1513,11 +1513,11 @@ impl<'a> Union<'a> {
     unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<DType>>>>(Union::VT_DTYPES, None)}
   }
   #[inline]
-  pub fn type_ids(&self) -> Option<::flatbuffers::Vector<'a, i8>> {
+  pub fn type_ids(&self) -> Option<::flatbuffers::Vector<'a, u8>> {
     // Safety:
     // Created from valid Table for this object
     // which contains a valid value in this slot
-    unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, i8>>>(Union::VT_TYPE_IDS, None)}
+    unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u8>>>(Union::VT_TYPE_IDS, None)}
   }
 }
 
@@ -1529,15 +1529,15 @@ impl ::flatbuffers::Verifiable for Union<'_> {
     v.visit_table(pos)?
      .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset<&'_ str>>>>("names", Self::VT_NAMES, false)?
      .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset<DType>>>>("dtypes", Self::VT_DTYPES, false)?
-     .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, i8>>>("type_ids", Self::VT_TYPE_IDS, false)?
+     .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u8>>>("type_ids", Self::VT_TYPE_IDS, false)?
      .finish();
     Ok(())
   }
 }
 pub struct UnionArgs<'a> {
     pub names: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>,
     pub dtypes: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<DType<'a>>>>>,
-    pub type_ids: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, i8>>>,
+    pub type_ids: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u8>>>,
 }
 impl<'a> Default for UnionArgs<'a> {
   #[inline]
@@ -1564,7 +1564,7 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> UnionBuilder<'a, 'b, A> {
     self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Union::VT_DTYPES, dtypes);
   }
   #[inline]
-  pub fn add_type_ids(&mut self, type_ids: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , i8>>) {
+  pub fn add_type_ids(&mut self, type_ids: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u8>>) {
     self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Union::VT_TYPE_IDS, type_ids);
   }
   #[inline]

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 {