Use UnionFields::from_fields

Author: alamb

datafusion/datasource-avro/src/avro_to_arrow/schema.rs

@@ -117,11 +117,8 @@ fn schema_to_field_with_props(
                     .iter()
                     .map(|s| schema_to_field_with_props(s, None, has_nullable, None))
                     .collect::<Result<Vec<Field>>>()?;
-                let type_ids = 0_i8..fields.len() as i8;
-                DataType::Union(
-                    UnionFields::try_new(type_ids, fields).unwrap(),
-                    UnionMode::Dense,
-                )
+                // Assign type_ids based on the order in which they appear
+                DataType::Union(UnionFields::from_fields(fields), UnionMode::Dense)
             }
         }
         AvroSchema::Record(RecordSchema { fields, .. }) => {

datafusion/proto-common/src/from_proto/mod.rs

@@ -304,16 +304,15 @@ impl TryFrom<&protobuf::arrow_type::ArrowTypeEnum> for DataType {
                 };
                 let union_fields = parse_proto_fields_to_fields(&union.union_types)?;
 
-                // Default to index based type ids if not provided
-                let type_ids: Vec<_> = match union.type_ids.is_empty() {
-                    true => (0..union_fields.len() as i8).collect(),
-                    false => union.type_ids.iter().map(|i| *i as i8).collect(),
-                };
-
-                let union_fields =
+                // Default to index based type ids if not explicitly provided
+                let union_fields = if union.type_ids.is_empty() {
+                    UnionFields::from_fields(union_fields)
+                } else {
+                    let type_ids = union.type_ids.iter().map(|i| *i as i8);
                     UnionFields::try_new(type_ids, union_fields).map_err(|e| {
                         DataFusionError::from(e).context("Deserializing Union DataType")
-                    })?;
+                    })?
+                };
                 DataType::Union(union_fields, union_mode)
             }
             arrow_type::ArrowTypeEnum::Dictionary(dict) => {