Change UnionArray constructors (apache#5623)

* Change `UnionArray` constructors

* Fix a comment

* Clippy and avoid using hashmaps

* Additional test

---------

Co-authored-by: Raphael Taylor-Davies <r.taylordavies@googlemail.com>

Author: mbrobbel

arrow-array/src/array/union_array.rs

@@ -23,7 +23,8 @@ use arrow_buffer::{ArrowNativeType, Buffer};
 use arrow_data::ArrayDataBuilder;
 use arrow_schema::{ArrowError, DataType, Field};
 use std::any::Any;
-use std::collections::HashMap;
+use std::collections::BTreeMap;
+use std::sync::Arc;
 
 /// `FieldData` is a helper struct to track the state of the fields in the `UnionBuilder`.
 #[derive(Debug)]
@@ -142,7 +143,7 @@ pub struct UnionBuilder {
     /// The current number of slots in the array
     len: usize,
     /// Maps field names to `FieldData` instances which track the builders for that field
-    fields: HashMap<String, FieldData>,
+    fields: BTreeMap<String, FieldData>,
     /// Builder to keep track of type ids
     type_id_builder: Int8BufferBuilder,
     /// Builder to keep track of offsets (`None` for sparse unions)
@@ -165,7 +166,7 @@ impl UnionBuilder {
     pub fn with_capacity_dense(capacity: usize) -> Self {
         Self {
             len: 0,
-            fields: HashMap::default(),
+            fields: Default::default(),
             type_id_builder: Int8BufferBuilder::new(capacity),
             value_offset_builder: Some(Int32BufferBuilder::new(capacity)),
             initial_capacity: capacity,
@@ -176,7 +177,7 @@ impl UnionBuilder {
     pub fn with_capacity_sparse(capacity: usize) -> Self {
         Self {
             len: 0,
-            fields: HashMap::default(),
+            fields: Default::default(),
             type_id_builder: Int8BufferBuilder::new(capacity),
             value_offset_builder: None,
             initial_capacity: capacity,
@@ -274,40 +275,39 @@ impl UnionBuilder {
     }
 
     /// Builds this builder creating a new `UnionArray`.
-    pub fn build(mut self) -> Result<UnionArray, ArrowError> {
-        let type_id_buffer = self.type_id_builder.finish();
-        let value_offsets_buffer = self.value_offset_builder.map(|mut b| b.finish());
-        let mut children = Vec::new();
-        for (
-            name,
-            FieldData {
-                type_id,
-                data_type,
-                mut values_buffer,
-                slots,
-                null_buffer_builder: mut bitmap_builder,
-            },
-        ) in self.fields.into_iter()
-        {
-            let buffer = values_buffer.finish();
-            let arr_data_builder = ArrayDataBuilder::new(data_type.clone())
-                .add_buffer(buffer)
-                .len(slots)
-                .nulls(bitmap_builder.finish());
-
-            let arr_data_ref = unsafe { arr_data_builder.build_unchecked() };
-            let array_ref = make_array(arr_data_ref);
-            children.push((type_id, (Field::new(name, data_type, false), array_ref)))
-        }
-
-        children.sort_by(|a, b| {
-            a.0.partial_cmp(&b.0)
-                .expect("This will never be None as type ids are always i8 values.")
-        });
-        let children: Vec<_> = children.into_iter().map(|(_, b)| b).collect();
-
-        let type_ids: Vec<i8> = (0_i8..children.len() as i8).collect();
-
-        UnionArray::try_new(&type_ids, type_id_buffer, value_offsets_buffer, children)
+    pub fn build(self) -> Result<UnionArray, ArrowError> {
+        let mut children = Vec::with_capacity(self.fields.len());
+        let union_fields = self
+            .fields
+            .into_iter()
+            .map(
+                |(
+                    name,
+                    FieldData {
+                        type_id,
+                        data_type,
+                        mut values_buffer,
+                        slots,
+                        mut null_buffer_builder,
+                    },
+                )| {
+                    let array_ref = make_array(unsafe {
+                        ArrayDataBuilder::new(data_type.clone())
+                            .add_buffer(values_buffer.finish())
+                            .len(slots)
+                            .nulls(null_buffer_builder.finish())
+                            .build_unchecked()
+                    });
+                    children.push(array_ref);
+                    (type_id, Arc::new(Field::new(name, data_type, false)))
+                },
+            )
+            .collect();
+        UnionArray::try_new(
+            union_fields,
+            self.type_id_builder.into(),
+            self.value_offset_builder.map(Into::into),
+            children,
+        )
     }
 }

arrow-array/src/builder/union_builder.rs

@@ -142,7 +142,7 @@ mod tests {
     use arrow_array::builder::*;
     use arrow_array::types::*;
     use arrow_array::*;
-    use arrow_buffer::Buffer;
+    use arrow_buffer::ScalarBuffer;
     use arrow_schema::*;
 
     use crate::display::array_value_to_string;
@@ -851,14 +851,18 @@ mod tests {
 
         // Can't use UnionBuilder with non-primitive types, so manually build outer UnionArray
         let a_array = Int32Array::from(vec![None, None, None, Some(1234), Some(23)]);
-        let type_ids = Buffer::from_slice_ref([1_i8, 1, 0, 0, 1]);
+        let type_ids = [1, 1, 0, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
 
-        let children: Vec<(Field, Arc<dyn Array>)> = vec![
-            (Field::new("a", DataType::Int32, true), Arc::new(a_array)),
-            (inner_field.clone(), Arc::new(inner)),
-        ];
+        let children = vec![Arc::new(a_array) as Arc<dyn Array>, Arc::new(inner)];
+
+        let union_fields = [
+            (0, Arc::new(Field::new("a", DataType::Int32, true))),
+            (1, Arc::new(inner_field.clone())),
+        ]
+        .into_iter()
+        .collect();
 
-        let outer = UnionArray::try_new(&[0, 1], type_ids, None, children).unwrap();
+        let outer = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
 
         let schema = Schema::new(vec![Field::new_union(
             "Teamsters",

arrow-cast/src/pretty.rs

@@ -597,20 +597,17 @@ fn hydrate_dictionary(array: &ArrayRef, data_type: &DataType) -> Result<ArrayRef
         (DataType::Union(_, UnionMode::Sparse), DataType::Union(fields, UnionMode::Sparse)) => {
             let union_arr = array.as_any().downcast_ref::<UnionArray>().unwrap();
 
-            let (type_ids, fields): (Vec<i8>, Vec<&FieldRef>) = fields.iter().unzip();
-
             Arc::new(UnionArray::try_new(
-                &type_ids,
-                union_arr.type_ids().inner().clone(),
+                fields.clone(),
+                union_arr.type_ids().clone(),
                 None,
                 fields
                     .iter()
-                    .enumerate()
-                    .map(|(col, field)| {
-                        Ok((
-                            field.as_ref().clone(),
-                            arrow_cast::cast(union_arr.child(col as i8), field.data_type())?,
-                        ))
+                    .map(|(type_id, field)| {
+                        Ok(arrow_cast::cast(
+                            union_arr.child(type_id),
+                            field.data_type(),
+                        )?)
                     })
                     .collect::<Result<Vec<_>>>()?,
             )?)
@@ -625,10 +622,10 @@ mod tests {
     use arrow_array::builder::StringDictionaryBuilder;
     use arrow_array::*;
     use arrow_array::{cast::downcast_array, types::*};
-    use arrow_buffer::Buffer;
+    use arrow_buffer::ScalarBuffer;
     use arrow_cast::pretty::pretty_format_batches;
     use arrow_ipc::MetadataVersion;
-    use arrow_schema::UnionMode;
+    use arrow_schema::{UnionFields, UnionMode};
     use std::collections::HashMap;
 
     use crate::decode::{DecodedPayload, FlightDataDecoder};
@@ -849,16 +846,23 @@ mod tests {
             true,
         )];
 
-        let type_ids = vec![0, 1, 2];
-        let union_fields = vec![
-            Field::new_list(
-                "dict_list",
-                Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
-                true,
+        let union_fields = [
+            (
+                0,
+                Arc::new(Field::new_list(
+                    "dict_list",
+                    Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
+                    true,
+                )),
             ),
-            Field::new_struct("struct", struct_fields.clone(), true),
-            Field::new("string", DataType::Utf8, true),
-        ];
+            (
+                1,
+                Arc::new(Field::new_struct("struct", struct_fields.clone(), true)),
+            ),
+            (2, Arc::new(Field::new("string", DataType::Utf8, true))),
+        ]
+        .into_iter()
+        .collect::<UnionFields>();
 
         let struct_fields = vec![Field::new_list(
             "dict_list",
@@ -872,21 +876,15 @@ mod tests {
 
         let arr1 = builder.finish();
 
-        let type_id_buffer = Buffer::from_slice_ref([0_i8]);
+        let type_id_buffer = [0].into_iter().collect::<ScalarBuffer<i8>>();
         let arr1 = UnionArray::try_new(
-            &type_ids,
+            union_fields.clone(),
             type_id_buffer,
             None,
             vec![
-                (union_fields[0].clone(), Arc::new(arr1)),
-                (
-                    union_fields[1].clone(),
-                    new_null_array(union_fields[1].data_type(), 1),
-                ),
-                (
-                    union_fields[2].clone(),
-                    new_null_array(union_fields[2].data_type(), 1),
-                ),
+                Arc::new(arr1) as Arc<dyn Array>,
+                new_null_array(union_fields.iter().nth(1).unwrap().1.data_type(), 1),
+                new_null_array(union_fields.iter().nth(2).unwrap().1.data_type(), 1),
             ],
         )
         .unwrap();
@@ -896,47 +894,36 @@ mod tests {
         let arr2 = Arc::new(builder.finish());
         let arr2 = StructArray::new(struct_fields.clone().into(), vec![arr2], None);
 
-        let type_id_buffer = Buffer::from_slice_ref([1_i8]);
+        let type_id_buffer = [1].into_iter().collect::<ScalarBuffer<i8>>();
         let arr2 = UnionArray::try_new(
-            &type_ids,
+            union_fields.clone(),
             type_id_buffer,
             None,
             vec![
-                (
-                    union_fields[0].clone(),
-                    new_null_array(union_fields[0].data_type(), 1),
-                ),
-                (union_fields[1].clone(), Arc::new(arr2)),
-                (
-                    union_fields[2].clone(),
-                    new_null_array(union_fields[2].data_type(), 1),
-                ),
+                new_null_array(union_fields.iter().next().unwrap().1.data_type(), 1),
+                Arc::new(arr2),
+                new_null_array(union_fields.iter().nth(2).unwrap().1.data_type(), 1),
             ],
         )
         .unwrap();
 
-        let type_id_buffer = Buffer::from_slice_ref([2_i8]);
+        let type_id_buffer = [2].into_iter().collect::<ScalarBuffer<i8>>();
         let arr3 = UnionArray::try_new(
-            &type_ids,
+            union_fields.clone(),
             type_id_buffer,
             None,
             vec![
-                (
-                    union_fields[0].clone(),
-                    new_null_array(union_fields[0].data_type(), 1),
-                ),
-                (
-                    union_fields[1].clone(),
-                    new_null_array(union_fields[1].data_type(), 1),
-                ),
-                (
-                    union_fields[2].clone(),
-                    Arc::new(StringArray::from(vec!["e"])),
-                ),
+                new_null_array(union_fields.iter().next().unwrap().1.data_type(), 1),
+                new_null_array(union_fields.iter().nth(1).unwrap().1.data_type(), 1),
+                Arc::new(StringArray::from(vec!["e"])),
             ],
         )
         .unwrap();
 
+        let (type_ids, union_fields): (Vec<_>, Vec<_>) = union_fields
+            .iter()
+            .map(|(type_id, field_ref)| (type_id, (*Arc::clone(field_ref)).clone()))
+            .unzip();
         let schema = Arc::new(Schema::new(vec![Field::new_union(
             "union",
             type_ids.clone(),

arrow-flight/src/encode.rs

@@ -21,6 +21,7 @@
 //!
 //! This is not a canonical format, but provides a human-readable way of verifying language implementations
 
+use arrow_buffer::ScalarBuffer;
 use hex::decode;
 use num::BigInt;
 use num::Signed;
@@ -835,26 +836,18 @@ pub fn array_from_json(
                 ));
             };
 
-            let offset: Option<Buffer> = json_col.offset.map(|offsets| {
-                let offsets: Vec<i32> =
-                    offsets.iter().map(|v| v.as_i64().unwrap() as i32).collect();
-                Buffer::from(&offsets.to_byte_slice())
-            });
+            let offset: Option<ScalarBuffer<i32>> = json_col
+                .offset
+                .map(|offsets| offsets.iter().map(|v| v.as_i64().unwrap() as i32).collect());
 
-            let mut children: Vec<(Field, Arc<dyn Array>)> = vec![];
+            let mut children = Vec::with_capacity(fields.len());
             for ((_, field), col) in fields.iter().zip(json_col.children.unwrap()) {
                 let array = array_from_json(field, col, dictionaries)?;
-                children.push((field.as_ref().clone(), array));
+                children.push(array);
             }
 
-            let field_type_ids = fields.iter().map(|(id, _)| id).collect::<Vec<_>>();
-            let array = UnionArray::try_new(
-                &field_type_ids,
-                Buffer::from(&type_ids.to_byte_slice()),
-                offset,
-                children,
-            )
-            .unwrap();
+            let array =
+                UnionArray::try_new(fields.clone(), type_ids.into(), offset, children).unwrap();
             Ok(Arc::new(array))
         }
         t => Err(ArrowError::JsonError(format!(

arrow-integration-test/src/lib.rs

@@ -31,7 +31,7 @@ use std::io::{BufReader, Read, Seek, SeekFrom};
 use std::sync::Arc;
 
 use arrow_array::*;
-use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer};
+use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ScalarBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::*;
 
@@ -214,26 +214,25 @@ fn create_array(
                 reader.next_buffer()?;
             }
 
-            let type_ids: Buffer = reader.next_buffer()?[..len].into();
+            let type_ids: ScalarBuffer<i8> = reader.next_buffer()?.slice_with_length(0, len).into();
 
             let value_offsets = match mode {
                 UnionMode::Dense => {
-                    let buffer = reader.next_buffer()?;
-                    Some(buffer[..len * 4].into())
+                    let offsets: ScalarBuffer<i32> =
+                        reader.next_buffer()?.slice_with_length(0, len * 4).into();
+                    Some(offsets)
                 }
                 UnionMode::Sparse => None,
             };
 
             let mut children = Vec::with_capacity(fields.len());
-            let mut ids = Vec::with_capacity(fields.len());
 
-            for (id, field) in fields.iter() {
+            for (_id, field) in fields.iter() {
                 let child = create_array(reader, field, variadic_counts, require_alignment)?;
-                children.push((field.as_ref().clone(), child));
-                ids.push(id);
+                children.push(child);
             }
 
-            let array = UnionArray::try_new(&ids, type_ids, value_offsets, children)?;
+            let array = UnionArray::try_new(fields.clone(), type_ids, value_offsets, children)?;
             Ok(Arc::new(array))
         }
         Null => {