RecordBatch normalization (flattening)

arrow-array/src/record_batch.rs

@@ -19,8 +19,8 @@
 //! [schema](arrow_schema::Schema).
 
 use crate::{new_empty_array, Array, ArrayRef, StructArray};
-use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaBuilder, SchemaRef};
-use std::ops::Index;
+use arrow_schema::{ArrowError, DataType, Field, Fields, Schema, SchemaBuilder, SchemaRef};
+use std::ops::{Deref, Index};
 use std::sync::Arc;
 
 /// Trait for types that can read `RecordBatch`'s.
@@ -1206,6 +1206,43 @@ mod tests {
         assert_ne!(batch1, batch2);
     }
 
+    #[test]
+    fn normalize() {
+        let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""]));
+        let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)]));
+        let year: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2022)]));
+
+        let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
+        let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
+        let year_field = Arc::new(Field::new("year", DataType::Int64, true));
+
+        let a = Arc::new(StructArray::from(vec![
+            (animals_field.clone(), Arc::new(animals) as ArrayRef),
+            (n_legs_field.clone(), Arc::new(n_legs) as ArrayRef),
+            (year_field.clone(), Arc::new(year) as ArrayRef),
+        ]));
+        let month = Arc::new(Int64Array::from(vec![Some(4), Some(6)]));
+
+        let schema = Schema::new(vec![
+            Field::new(
+                "a",
+                DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])),
+                false,
+            ),
+            Field::new("month", DataType::Int64, true),
+        ]);
+        let normalized = schema.clone().normalize(".", 0).unwrap();
+        println!("{:?}", normalized);
+
+        let record_batch =
+            RecordBatch::try_new(Arc::new(schema), vec![a, month]).expect("valid conversion");
+
+        println!("Fields: {:?}", record_batch.schema().fields());
+        println!("Metadata{:?}", record_batch.columns());
+
+        //println!("{:?}", record_batch);
+    }
+
     #[test]
     fn project() {
         let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]));
@@ -1318,7 +1355,9 @@ mod tests {
         let metadata = vec![("foo".to_string(), "bar".to_string())]
             .into_iter()
             .collect();
+        println!("Metadata: {:?}", metadata);
         let metadata_schema = nullable_schema.as_ref().clone().with_metadata(metadata);
+        println!("Metadata schema: {:?}", metadata_schema);
         let batch = batch.with_schema(Arc::new(metadata_schema)).unwrap();
 
         // Cannot remove metadata

arrow-schema/src/schema.rs

@@ -22,7 +22,7 @@ use std::sync::Arc;
 
 use crate::error::ArrowError;
 use crate::field::Field;
-use crate::{FieldRef, Fields};
+use crate::{DataType, FieldRef, Fields};
 
 /// A builder to facilitate building a [`Schema`] from iteratively from [`FieldRef`]
 #[derive(Debug, Default)]
@@ -413,6 +413,91 @@ impl Schema {
         &self.metadata
     }
 
+    pub fn normalize(self, separator: &str, mut max_level: usize) -> Result<Self, ArrowError> {
+        if max_level == 0 {
+            max_level = usize::MAX;
+        }
+        let mut new_fields: Vec<Field> = vec![];
+        for field in self.fields() {
+            match field.data_type() {
+                //DataType::List(f) => field,
+                //DataType::ListView(_) => field,
+                //DataType::FixedSizeList(_, _) => field,
+                //DataType::LargeList(_) => field,
+                //DataType::LargeListView(_) => field,
+                DataType::Struct(nested_fields) => {
+                    let field_name = field.name().as_str();
+                    new_fields = [
+                        new_fields,
+                        Self::normalizer(
+                            nested_fields.to_vec(),
+                            field_name,
+                            separator,
+                            max_level - 1,
+                        ),
+                    ]
+                    .concat();
+                }
+                //DataType::Union(_, _) => field,
+                //DataType::Dictionary(_, _) => field,
+                //DataType::Map(_, _) => field,
+                //DataType::RunEndEncoded(_, _) => field, // not sure how to support this field
+                _ => new_fields.push(Field::new(
+                    field.name(),
+                    field.data_type().clone(),
+                    field.is_nullable(),
+                )),
+            };
+        }
+        Ok(Self::new_with_metadata(new_fields, self.metadata.clone()))
+    }
+
+    fn normalizer(
+        fields: Vec<FieldRef>,
+        key_string: &str,
+        separator: &str,
+        max_level: usize,
+    ) -> Vec<Field> {
+        if max_level > 0 {
+            let mut new_fields: Vec<Field> = vec![];
+            for field in fields {
+                match field.data_type() {
+                    //DataType::List(f) => field,
+                    //DataType::ListView(_) => field,
+                    //DataType::FixedSizeList(_, _) => field,
+                    //DataType::LargeList(_) => field,
+                    //DataType::LargeListView(_) => field,
+                    DataType::Struct(nested_fields) => {
+                        let field_name = field.name().as_str();
+                        let new_key = format!("{key_string}{separator}{field_name}");
+                        new_fields = [
+                            new_fields,
+                            Self::normalizer(
+                                nested_fields.to_vec(),
+                                new_key.as_str(),
+                                separator,
+                                max_level - 1,
+                            ),
+                        ]
+                        .concat();
+                    }
+                    //DataType::Union(_, _) => field,
+                    //DataType::Dictionary(_, _) => field,
+                    //DataType::Map(_, _) => field,
+                    //DataType::RunEndEncoded(_, _) => field, // not sure how to support this field
+                    _ => new_fields.push(Field::new(
+                        format!("{key_string}{separator}{}", field.name()),
+                        field.data_type().clone(),
+                        field.is_nullable(),
+                    )),
+                };
+            }
+            new_fields
+        } else {
+            todo!()
+        }
+    }
+
     /// Look up a column by name and return a immutable reference to the column along with
     /// its index.
     pub fn column_with_name(&self, name: &str) -> Option<(usize, &Field)> {