[arrow-rs] branch master updated: Add RecordBatch::with_schema (#4028)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new e45f8b513 Add RecordBatch::with_schema (#4028)
e45f8b513 is described below

commit e45f8b5132c10505a6f2b9d8a8e8e9207049605f
Author: Raphael Taylor-Davies <1781103+tustv...@users.noreply.github.com>
AuthorDate: Fri Apr 7 09:20:10 2023 +0100

    Add RecordBatch::with_schema (#4028)
---
 arrow-array/src/record_batch.rs | 49 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs
index 081bd55fc..1350285f8 100644
--- a/arrow-array/src/record_batch.rs
+++ b/arrow-array/src/record_batch.rs
@@ -204,6 +204,25 @@ impl RecordBatch {
         })
     }
 
+    /// Override the schema of this [`RecordBatch`]
+    ///
+    /// Returns an error if `schema` is not a superset of the current schema
+    /// as determined by [`Schema::contains`]
+    pub fn with_schema(self, schema: SchemaRef) -> Result<Self, ArrowError> {
+        if !schema.contains(self.schema.as_ref()) {
+            return Err(ArrowError::SchemaError(format!(
+                "{schema} is not a superset of {}",
+                self.schema
+            )));
+        }
+
+        Ok(Self {
+            schema,
+            columns: self.columns,
+            row_count: self.row_count,
+        })
+    }
+
     /// Returns the [`Schema`](arrow_schema::Schema) of the record batch.
     pub fn schema(&self) -> SchemaRef {
         self.schema.clone()
@@ -1062,4 +1081,34 @@ mod tests {
         ));
         let _ = RecordBatch::from(s);
     }
+
+    #[test]
+    fn test_with_schema() {
+        let required_schema = Schema::new(vec![Field::new("a", 
DataType::Int32, false)]);
+        let required_schema = Arc::new(required_schema);
+        let nullable_schema = Schema::new(vec![Field::new("a", 
DataType::Int32, true)]);
+        let nullable_schema = Arc::new(nullable_schema);
+
+        let batch = RecordBatch::try_new(
+            required_schema.clone(),
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as _],
+        )
+        .unwrap();
+
+        // Can add nullability
+        let batch = batch.with_schema(nullable_schema.clone()).unwrap();
+
+        // Cannot remove nullability
+        batch.clone().with_schema(required_schema).unwrap_err();
+
+        // Can add metadata
+        let metadata = vec![("foo".to_string(), "bar".to_string())]
+            .into_iter()
+            .collect();
+        let metadata_schema = 
nullable_schema.as_ref().clone().with_metadata(metadata);
+        let batch = batch.with_schema(Arc::new(metadata_schema)).unwrap();
+
+        // Cannot remove metadata
+        batch.with_schema(nullable_schema).unwrap_err();
+    }
 }