Don't omit schema metadata when removing column (#5328)

* Don't omit schema metadata when removing column

* Add test

* Update arrow-schema/src/schema.rs

Co-authored-by: Andrew Lamb <[email protected]>

---------

Co-authored-by: Andrew Lamb <[email protected]>

arrow-array/src/record_batch.rs

@@ -355,7 +355,7 @@ impl RecordBatch {
     /// assert_eq!(batch.num_columns(), 1);
     /// ```
     pub fn remove_column(&mut self, index: usize) -> ArrayRef {
-        let mut builder = SchemaBuilder::from(self.schema.fields());
+        let mut builder = SchemaBuilder::from(self.schema.as_ref());
         builder.remove(index);
         self.schema = Arc::new(builder.finish());
         self.columns.remove(index)
@@ -618,6 +618,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use super::*;
     use crate::{BooleanArray, Int32Array, Int64Array, Int8Array, ListArray, StringArray};
     use arrow_buffer::{Buffer, ToByteSlice};
@@ -1155,4 +1157,31 @@ mod tests {
         let size = get_size(reader);
         assert_eq!(size, 0);
     }
+
+    #[test]
+    fn test_remove_column_maintains_schema_metadata() {
+        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
+        let bool_array = BooleanArray::from(vec![true, false, false, true, true]);
+
+        let mut metadata = HashMap::new();
+        metadata.insert("foo".to_string(), "bar".to_string());
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("bool", DataType::Boolean, false),
+        ])
+        .with_metadata(metadata);
+
+        let mut batch = RecordBatch::try_new(
+            Arc::new(schema),
+            vec![Arc::new(id_array), Arc::new(bool_array)],
+        )
+        .unwrap();
+
+        let _removed_column = batch.remove_column(0);
+        assert_eq!(batch.schema().metadata().len(), 1);
+        assert_eq!(
+            batch.schema().metadata().get("foo").unwrap().as_str(),
+            "bar"
+        );
+    }
 }

arrow-schema/src/schema.rs

@@ -140,6 +140,12 @@ impl From<Fields> for SchemaBuilder {
     }
 }
 
+impl From<&Schema> for SchemaBuilder {
+    fn from(value: &Schema) -> Self {
+        Self::from(value.clone())
+    }
+}
+
 impl From<Schema> for SchemaBuilder {
     fn from(value: Schema) -> Self {
         Self {