StructType modification method and stat_transform schema boilerplate code refactor. (delta-io#1872)

## What changes are proposed in this pull request?

Closes delta-io#1657

Currently, there is no Schema modification method supporting
insert/remove/replace a field from the Schema. Schema modification is
left to the caller to construct, which create the boilerplate code
across the codebase. This PR introduced the schema modification method
for inserting, removing and replacing a field from schema. In addition,
the schema modification method is integrated into the stats_transform to
reduce boilerplate code when creating schema.

Change included:
- New StructType method for schema modification : 
- `pub fn with_field_inserted_after(mut self, after: Option<&str>,
new_field: StructField) -> Self`
    -  `pub fn with_field_removed(mut self, name: &str) -> Self`
- `pub fn with_field_replaced(mut self, name: &str, new_field:
StructField)`
- `pub fn with_field_inserted_after(mut self, before: Option<&str>,
new_field: StructField) -> Self`
- Unit-test for StructType to verify the correct behaviour of the schema
modification method `schema::tests::test_with_field_`
- Update the stats_transform.rs function `transform_add_schema`,
`build_add_output_schema`, `add_stats_parsed_to_add_schema`

## How was this change tested?
- Unit test for Schema modification method: `cargo test --package
delta_kernel --lib -- schema::tests::test_with_field_`
- Unit test for stat_transform `cargo test --package delta_kernel --lib
-- checkpoint::stats_transform::test`

---------

Co-authored-by: Nick Lanham <nicklan@users.noreply.github.com>

Author: thinh2

kernel/src/checkpoint/stats_transform.rs

@@ -121,7 +121,7 @@ pub(crate) fn build_checkpoint_read_schema_with_stats(
                 "stats_parsed field already exists in Add schema",
             ));
         }
-        Ok(add_stats_parsed_to_add_schema(add_struct, stats_schema))
+        add_stats_parsed_to_add_schema(add_struct, stats_schema)
     })
 }
 
@@ -140,7 +140,7 @@ pub(crate) fn build_checkpoint_output_schema(
     stats_schema: &StructType,
 ) -> DeltaResult<SchemaRef> {
     transform_add_schema(base_schema, |add_struct| {
-        Ok(build_add_output_schema(config, add_struct, stats_schema))
+        build_add_output_schema(config, add_struct, stats_schema)
     })
 }
 
@@ -211,75 +211,55 @@ fn transform_add_schema(
     };
 
     let modified_add = transform_fn(add_struct)?;
-    let fields: Vec<StructField> = base_schema
-        .fields()
-        .map(|field| {
-            if field.name == ADD_NAME {
-                StructField {
-                    name: field.name.clone(),
-                    data_type: DataType::Struct(Box::new(modified_add.clone())),
-                    nullable: field.nullable,
-                    metadata: field.metadata.clone(),
-                }
-            } else {
-                field.clone()
-            }
-        })
-        .collect();
-
-    Ok(Arc::new(StructType::new_unchecked(fields)))
+    let new_schema = base_schema.clone().with_field_replaced(
+        ADD_NAME,
+        StructField {
+            name: ADD_NAME.to_string(),
+            data_type: DataType::Struct(Box::new(modified_add)),
+            nullable: add_field.nullable,
+            metadata: add_field.metadata.clone(),
+        },
+    )?;
+
+    Ok(Arc::new(new_schema))
 }
 
 /// Adds `stats_parsed` field after `stats` in the Add action schema.
 fn add_stats_parsed_to_add_schema(
     add_schema: &StructType,
     stats_schema: &StructType,
-) -> StructType {
-    let mut fields: Vec<StructField> = Vec::with_capacity(add_schema.num_fields() + 1);
-
-    for field in add_schema.fields() {
-        fields.push(field.clone());
-        if field.name == STATS_FIELD {
-            // Insert stats_parsed right after stats
-            fields.push(StructField::nullable(
-                STATS_PARSED_FIELD,
-                DataType::Struct(Box::new(stats_schema.clone())),
-            ));
-        }
-    }
-
-    StructType::new_unchecked(fields)
+) -> DeltaResult<StructType> {
+    let new_schema = add_schema.clone();
+    new_schema.with_field_inserted_after(
+        Some(STATS_FIELD),
+        StructField::nullable(
+            STATS_PARSED_FIELD,
+            DataType::Struct(Box::new(stats_schema.clone())),
+        ),
+    )
 }
 
 fn build_add_output_schema(
     config: &StatsTransformConfig,
     add_schema: &StructType,
     stats_schema: &StructType,
-) -> StructType {
-    let capacity = add_schema.num_fields()
-        - if config.write_stats_as_json { 0 } else { 1 } // dropping stats?
-        + if config.write_stats_as_struct { 1 } else { 0 }; // adding stats_parsed?
-    let mut fields: Vec<StructField> = Vec::with_capacity(capacity);
-
-    for field in add_schema.fields() {
-        if field.name == STATS_FIELD {
-            // Include stats if writing as JSON
-            if config.write_stats_as_json {
-                fields.push(field.clone());
-            }
-            // Add stats_parsed after stats position if writing as struct
-            if config.write_stats_as_struct {
-                fields.push(StructField::nullable(
-                    STATS_PARSED_FIELD,
-                    DataType::Struct(Box::new(stats_schema.clone())),
-                ));
-            }
-        } else {
-            fields.push(field.clone());
-        }
+) -> DeltaResult<StructType> {
+    let mut new_schema = add_schema.clone();
+    if config.write_stats_as_struct {
+        new_schema = new_schema.with_field_inserted_after(
+            Some(STATS_FIELD),
+            StructField::nullable(
+                STATS_PARSED_FIELD,
+                DataType::Struct(Box::new(stats_schema.clone())),
+            ),
+        )?
     }
 
-    StructType::new_unchecked(fields)
+    if config.write_stats_as_json {
+        Ok(new_schema)
+    } else {
+        Ok(new_schema.with_field_removed(STATS_FIELD))
+    }
 }
 
 #[cfg(test)]
@@ -469,7 +449,8 @@ mod tests {
         let stats_schema =
             StructType::new_unchecked([StructField::nullable("numRecords", DataType::LONG)]);
 
-        let result = add_stats_parsed_to_add_schema(&add_schema, &stats_schema);
+        let result = add_stats_parsed_to_add_schema(&add_schema, &stats_schema)
+            .expect("add stats_parsed to add schema should produce a valid schema");
 
         // Should have 4 fields: path, stats, stats_parsed, tags
         assert_eq!(result.fields().count(), 4);
@@ -492,7 +473,8 @@ mod tests {
 
         let stats_schema = StructType::new_unchecked([]);
 
-        let result = build_add_output_schema(&config, &add_schema, &stats_schema);
+        let result = build_add_output_schema(&config, &add_schema, &stats_schema)
+            .expect("build add output schema should produce a valid schema");
 
         // Should have path and stats, no stats_parsed
         let field_names: Vec<&str> = result.fields().map(|f| f.name.as_str()).collect();
@@ -514,7 +496,8 @@ mod tests {
         let stats_schema =
             StructType::new_unchecked([StructField::nullable("numRecords", DataType::LONG)]);
 
-        let result = build_add_output_schema(&config, &add_schema, &stats_schema);
+        let result = build_add_output_schema(&config, &add_schema, &stats_schema)
+            .expect("build add output schema should produce a valid schema");
 
         // Should have path and stats_parsed (stats dropped)
         let field_names: Vec<&str> = result.fields().map(|f| f.name.as_str()).collect();
@@ -536,7 +519,8 @@ mod tests {
         let stats_schema =
             StructType::new_unchecked([StructField::nullable("numRecords", DataType::LONG)]);
 
-        let result = build_add_output_schema(&config, &add_schema, &stats_schema);
+        let result = build_add_output_schema(&config, &add_schema, &stats_schema)
+            .expect("build add output schema should produce a valid schema");
 
         // Should have path, stats, and stats_parsed
         let field_names: Vec<&str> = result.fields().map(|f| f.name.as_str()).collect();