diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index 2e93acb0931c..557663922121 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -323,7 +323,7 @@ fn prepare_schema_for_flight(schema: &Schema) -> Schema { }) .collect(); - Schema::new(fields) + Schema::new(fields).with_metadata(schema.metadata().clone()) } /// Split [`RecordBatch`] so it hopefully fits into a gRPC response. @@ -461,6 +461,7 @@ mod tests { use arrow_array::{ DictionaryArray, Int16Array, Int32Array, Int64Array, StringArray, UInt64Array, }; + use std::collections::HashMap; use super::*; @@ -502,6 +503,17 @@ mod tests { ); } + #[test] + fn test_schema_metadata_encoded() { + let schema = + Schema::new(vec![Field::new("data", DataType::Int32, false)]).with_metadata( + HashMap::from([("some_key".to_owned(), "some_value".to_owned())]), + ); + + let got = prepare_schema_for_flight(&schema); + assert!(got.metadata().contains_key("some_key")); + } + #[test] fn test_encode_no_column_batch() { let batch = RecordBatch::try_new_with_options( diff --git a/arrow-flight/tests/encode_decode.rs b/arrow-flight/tests/encode_decode.rs index 25e74cb3b6bc..8c73a516b2b0 100644 --- a/arrow-flight/tests/encode_decode.rs +++ b/arrow-flight/tests/encode_decode.rs @@ -17,7 +17,7 @@ //! Tests for round trip encoding / decoding -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use arrow::{compute::concat_batches, datatypes::Int32Type}; use arrow_array::{ArrayRef, DictionaryArray, Float64Array, RecordBatch, UInt8Array}; @@ -62,6 +62,18 @@ async fn test_primative_one() { roundtrip(vec![make_primative_batch(5)]).await; } +#[tokio::test] +async fn test_schema_metadata() { + let batch = make_primative_batch(5); + let metadata = HashMap::from([("some_key".to_owned(), "some_value".to_owned())]); + + // create a batch that has schema level metadata + let schema = Arc::new(batch.schema().as_ref().clone().with_metadata(metadata)); + let batch = RecordBatch::try_new(schema, batch.columns().to_vec()).unwrap(); + + roundtrip(vec![batch]).await; +} + #[tokio::test] async fn test_primative_many() { roundtrip(vec![