fix schema/array column ordering issue by using BTreeMaps in the Trac…

…er and Builder
chmp · Oct 18, 2024 · 92804e9 · 92804e9
1 parent 87a0c79
commit 92804e9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/serde_arrow/src/internal/schema/tracer.rs b/serde_arrow/src/internal/schema/tracer.rs
@@ -1083,7 +1083,7 @@ impl UnionTracer {
                 STRATEGY_KEY.to_string(),
                 Strategy::EnumsWithNamedFieldsAsStructs.to_string(),
             );
-            let mut fields = Vec::new();
+            let mut fields = BTreeMap::new();
 
             // For this option, we want to merge the variant children up one level, combining the names
             // For each variant with name variant_name
@@ -1094,14 +1094,16 @@ impl UnionTracer {
                 if let Some(variant) = variant {
                     let schema = variant.tracer.to_schema()?;
                     for field in schema.fields {
-                        fields.push(field.to_flattened_union_field(variant.name.as_str()))
+                        let flat_field = field.to_flattened_union_field(variant.name.as_str());
+                        fields.insert(flat_field.name.to_string(), flat_field);
                     }
                 } else {
-                    fields.push(unknown_variant_field())
+                    let uf = unknown_variant_field();
+                    fields.insert(uf.name, unknown_variant_field());
                 };
             }
 
-            data_type = DataType::Struct(fields);
+            data_type = DataType::Struct(fields.into_values().collect());
         } else {
             let mut fields = Vec::new();
             for (idx, variant) in self.variants.iter().enumerate() {

diff --git a/serde_arrow/src/internal/serialization/flattened_union_builder.rs b/serde_arrow/src/internal/serialization/flattened_union_builder.rs
@@ -40,7 +40,7 @@ impl FlattenedUnionBuilder {
     }
 
     pub fn into_array(self) -> Result<Array> {
-        let mut fields = Vec::new();
+        let mut fields = BTreeMap::new();
 
         for (builder, meta) in self.fields.into_iter() {
             let ArrayBuilder::Struct(builder) = builder else {
@@ -53,13 +53,16 @@ impl FlattenedUnionBuilder {
                 // Name change is currently needed for struct field lookup to work correctly.
 
                 sub_meta.name = format!("{}::{}", meta.name, sub_meta.name);
-                fields.push((sub_builder.into_array()?, sub_meta));
+                fields.insert(
+                    sub_meta.name.to_owned(),
+                    (sub_builder.into_array()?, sub_meta),
+                );
             }
         }
 
         Ok(Array::Struct(StructArray {
             len: self.row_count,
-            fields,
+            fields: fields.into_values().collect(),
             // TODO: is this ok to hardcode?
             // assuming so because when testing manually,
             // validity of struct with nullable fields was None