feat: add custom_metadata support to RecordBatch with IPC read/write

Author: rustyconover

arrow-array/src/record_batch.rs

@@ -21,6 +21,7 @@
 use crate::cast::AsArray;
 use crate::{Array, ArrayRef, StructArray, new_empty_array};
 use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaBuilder, SchemaRef};
+use std::collections::HashMap;
 use std::ops::Index;
 use std::sync::Arc;
 
@@ -229,6 +230,12 @@ pub struct RecordBatch {
     ///
     /// This is stored separately from the columns to handle the case of no columns
     row_count: usize,
+
+    /// Per-batch custom metadata
+    ///
+    /// This corresponds to the `custom_metadata` field on the IPC `Message`
+    /// flatbuffer, allowing per-batch metadata separate from schema-level metadata.
+    custom_metadata: HashMap<String, String>,
 }
 
 impl RecordBatch {
@@ -289,6 +296,7 @@ impl RecordBatch {
             schema,
             columns,
             row_count,
+            custom_metadata: HashMap::new(),
         }
     }
 
@@ -316,6 +324,7 @@ impl RecordBatch {
             schema,
             columns,
             row_count: 0,
+            custom_metadata: HashMap::new(),
         }
     }
 
@@ -390,14 +399,30 @@ impl RecordBatch {
             schema,
             columns,
             row_count,
+            custom_metadata: HashMap::new(),
         })
     }
 
     /// Return the schema, columns and row count of this [`RecordBatch`]
+    ///
+    /// Note: this discards any [`Self::custom_metadata`]. Use
+    /// [`Self::into_parts_with_custom_metadata`] to also retrieve it.
     pub fn into_parts(self) -> (SchemaRef, Vec<ArrayRef>, usize) {
         (self.schema, self.columns, self.row_count)
     }
 
+    /// Return the schema, columns, row count and custom metadata of this [`RecordBatch`]
+    pub fn into_parts_with_custom_metadata(
+        self,
+    ) -> (SchemaRef, Vec<ArrayRef>, usize, HashMap<String, String>) {
+        (
+            self.schema,
+            self.columns,
+            self.row_count,
+            self.custom_metadata,
+        )
+    }
+
     /// Override the schema of this [`RecordBatch`]
     ///
     /// Returns an error if `schema` is not a superset of the current schema
@@ -416,6 +441,7 @@ impl RecordBatch {
             schema,
             columns: self.columns,
             row_count: self.row_count,
+            custom_metadata: self.custom_metadata,
         })
     }
 
@@ -451,6 +477,25 @@ impl RecordBatch {
         &mut schema.metadata
     }
 
+    /// Returns a reference to the per-batch custom metadata.
+    ///
+    /// This metadata corresponds to the `custom_metadata` field on the IPC
+    /// `Message` flatbuffer, separate from schema-level metadata.
+    pub fn custom_metadata(&self) -> &HashMap<String, String> {
+        &self.custom_metadata
+    }
+
+    /// Returns a mutable reference to the per-batch custom metadata.
+    pub fn custom_metadata_mut(&mut self) -> &mut HashMap<String, String> {
+        &mut self.custom_metadata
+    }
+
+    /// Sets the per-batch custom metadata, returning `self`.
+    pub fn with_custom_metadata(mut self, metadata: HashMap<String, String>) -> Self {
+        self.custom_metadata = metadata;
+        self
+    }
+
     /// Projects the schema onto the specified columns
     pub fn project(&self, indices: &[usize]) -> Result<RecordBatch, ArrowError> {
         let projected_schema = self.schema.project(indices)?;
@@ -475,7 +520,8 @@ impl RecordBatch {
                 SchemaRef::new(projected_schema),
                 batch_fields,
                 self.row_count,
-            ))
+            )
+            .with_custom_metadata(self.custom_metadata.clone()))
         }
     }
 
@@ -570,7 +616,9 @@ impl RecordBatch {
                 }
             }
         }
+        let custom_metadata = self.custom_metadata.clone();
         RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+            .map(|b| b.with_custom_metadata(custom_metadata))
     }
 
     /// Returns the number of columns in the record batch.
@@ -691,6 +739,7 @@ impl RecordBatch {
             schema: self.schema.clone(),
             columns,
             row_count: length,
+            custom_metadata: self.custom_metadata.clone(),
         }
     }
 
@@ -864,6 +913,7 @@ impl From<StructArray> for RecordBatch {
             schema: Arc::new(Schema::new(fields)),
             row_count,
             columns,
+            custom_metadata: HashMap::new(),
         }
     }
 }
@@ -1792,4 +1842,81 @@ mod tests {
         assert!(col.is_null(1));
         assert!(col.is_valid(2));
     }
+
+    #[test]
+    fn test_with_custom_metadata() {
+        let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+        assert!(batch.custom_metadata().is_empty());
+
+        let mut metadata = HashMap::new();
+        metadata.insert("key".to_string(), "value".to_string());
+        let batch = batch.with_custom_metadata(metadata.clone());
+        assert_eq!(batch.custom_metadata(), &metadata);
+    }
+
+    #[test]
+    fn test_custom_metadata_mut() {
+        let mut batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+        batch
+            .custom_metadata_mut()
+            .insert("key".to_string(), "value".to_string());
+        assert_eq!(
+            batch.custom_metadata().get("key"),
+            Some(&"value".to_string())
+        );
+    }
+
+    #[test]
+    fn test_slice_preserves_custom_metadata() {
+        let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+        let mut metadata = HashMap::new();
+        metadata.insert("k".to_string(), "v".to_string());
+        let batch = batch.with_custom_metadata(metadata.clone());
+
+        let sliced = batch.slice(0, 2);
+        assert_eq!(sliced.custom_metadata(), &metadata);
+    }
+
+    #[test]
+    fn test_project_preserves_custom_metadata() {
+        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
+        let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap();
+
+        let mut metadata = HashMap::new();
+        metadata.insert("k".to_string(), "v".to_string());
+        let batch = batch.with_custom_metadata(metadata.clone());
+
+        let projected = batch.project(&[0]).unwrap();
+        assert_eq!(projected.custom_metadata(), &metadata);
+    }
+
+    #[test]
+    fn test_into_parts_with_custom_metadata() {
+        let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+        let mut metadata = HashMap::new();
+        metadata.insert("k".to_string(), "v".to_string());
+        let batch = batch.with_custom_metadata(metadata.clone());
+
+        let (schema, columns, row_count, custom_metadata) = batch.into_parts_with_custom_metadata();
+        assert_eq!(schema.fields().len(), 1);
+        assert_eq!(columns.len(), 1);
+        assert_eq!(row_count, 3);
+        assert_eq!(custom_metadata, metadata);
+    }
+
+    #[test]
+    fn test_custom_metadata_equality() {
+        let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+        let batch2 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+
+        // Both empty metadata -> equal
+        assert_eq!(batch1, batch2);
+
+        // Different metadata -> not equal
+        let mut metadata = HashMap::new();
+        metadata.insert("k".to_string(), "v".to_string());
+        let batch1 = batch1.with_custom_metadata(metadata);
+        assert_ne!(batch1, batch2);
+    }
 }

arrow-flight/src/utils.rs

@@ -61,6 +61,8 @@ pub fn flight_data_to_arrow_batch(
     let message = arrow_ipc::root_as_message(&data.data_header[..])
         .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))?;
 
+    let custom_metadata = arrow_ipc::reader::message_custom_metadata(&message);
+
     message
         .header_as_record_batch()
         .ok_or_else(|| {
@@ -77,6 +79,13 @@ pub fn flight_data_to_arrow_batch(
                 None,
                 &message.version(),
             )
+            .map(|rb| {
+                if custom_metadata.is_empty() {
+                    rb
+                } else {
+                    rb.with_custom_metadata(custom_metadata)
+                }
+            })
         })?
 }
 

arrow-ipc/src/reader.rs

@@ -47,6 +47,21 @@ use crate::r#gen::Message::{self};
 use crate::{Block, CONTINUATION_MARKER, FieldNode, MetadataVersion};
 use DataType::*;
 
+/// Extract `custom_metadata` key-value pairs from an IPC [`Message`].
+///
+/// Returns an empty [`HashMap`] if the message has no custom metadata.
+pub fn message_custom_metadata(message: &crate::Message) -> HashMap<String, String> {
+    let mut metadata = HashMap::new();
+    if let Some(list) = message.custom_metadata() {
+        for kv in list {
+            if let (Some(k), Some(v)) = (kv.key(), kv.value()) {
+                metadata.insert(k.to_string(), v.to_string());
+            }
+        }
+    }
+    metadata
+}
+
 /// Read a buffer based on offset and length
 /// From <https://github.com/apache/arrow/blob/6a936c4ff5007045e86f65f1a6b6c3c955ad5103/format/Message.fbs#L58>
 /// Each constituent buffer is first compressed with the indicated
@@ -470,6 +485,8 @@ pub struct RecordBatchDecoder<'a> {
     ///
     /// See [`FileDecoder::with_skip_validation`] for details.
     skip_validation: UnsafeFlag,
+    /// Per-batch custom metadata to attach to the decoded RecordBatch
+    custom_metadata: HashMap<String, String>,
 }
 
 impl<'a> RecordBatchDecoder<'a> {
@@ -506,6 +523,7 @@ impl<'a> RecordBatchDecoder<'a> {
             projection: None,
             require_alignment: false,
             skip_validation: UnsafeFlag::new(),
+            custom_metadata: HashMap::new(),
         })
     }
 
@@ -544,6 +562,12 @@ impl<'a> RecordBatchDecoder<'a> {
         self
     }
 
+    /// Set per-batch custom metadata to attach to the decoded [`RecordBatch`]
+    pub(crate) fn with_custom_metadata(mut self, custom_metadata: HashMap<String, String>) -> Self {
+        self.custom_metadata = custom_metadata;
+        self
+    }
+
     /// Read the record batch, consuming the reader
     fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
         let mut variadic_counts: VecDeque<i64> = self
@@ -554,9 +578,10 @@ impl<'a> RecordBatchDecoder<'a> {
             .collect();
 
         let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
+        let custom_metadata = std::mem::take(&mut self.custom_metadata);
 
         let schema = Arc::clone(&self.schema);
-        if let Some(projection) = self.projection {
+        let batch = if let Some(projection) = self.projection {
             let mut arrays = vec![];
             // project fields
             for (idx, field) in schema.fields().iter().enumerate() {
@@ -608,7 +633,15 @@ impl<'a> RecordBatchDecoder<'a> {
                 assert!(variadic_counts.is_empty());
                 RecordBatch::try_new_with_options(schema, children, &options)
             }
-        }
+        };
+
+        batch.map(|b| {
+            if custom_metadata.is_empty() {
+                b
+            } else {
+                b.with_custom_metadata(custom_metadata)
+            }
+        })
     }
 
     fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
@@ -752,6 +785,11 @@ impl<'a> RecordBatchDecoder<'a> {
 /// and copy over the data if any array data in the input `buf` is not properly aligned.
 /// (Properly aligned array data will remain zero-copy.)
 /// Under the hood it will use [`arrow_data::ArrayDataBuilder::build_aligned`] to construct [`arrow_data::ArrayData`].
+///
+/// Note: this function operates on the inner `RecordBatch` flatbuffer, not the
+/// outer `Message` envelope. Message-level `custom_metadata` is not extracted.
+/// Callers who need it should use [`message_custom_metadata`] on the `Message`
+/// and apply it via [`RecordBatch::with_custom_metadata`].
 pub fn read_record_batch(
     buf: &Buffer,
     batch: crate::RecordBatch,
@@ -1114,6 +1152,7 @@ impl FileDecoder {
                 let batch = message.header_as_record_batch().ok_or_else(|| {
                     ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
                 })?;
+                let custom_metadata = message_custom_metadata(&message);
                 // read the block that makes up the record batch into a buffer
                 RecordBatchDecoder::try_new(
                     &buf.slice(block.metaDataLength() as _),
@@ -1125,6 +1164,7 @@ impl FileDecoder {
                 .with_projection(self.projection.as_deref())
                 .with_require_alignment(self.require_alignment)
                 .with_skip_validation(self.skip_validation.clone())
+                .with_custom_metadata(custom_metadata)
                 .read_record_batch()
                 .map(Some)
             }
@@ -1679,6 +1719,7 @@ impl<R: Read> StreamReader<R> {
                     ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
                 })?;
 
+                let custom_metadata = message_custom_metadata(&message);
                 let version = message.version();
                 let schema = self.schema.clone();
                 let record_batch = RecordBatchDecoder::try_new(
@@ -1691,6 +1732,7 @@ impl<R: Read> StreamReader<R> {
                 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
                 .with_require_alignment(false)
                 .with_skip_validation(self.skip_validation.clone())
+                .with_custom_metadata(custom_metadata)
                 .read_record_batch()?;
                 IpcMessage::RecordBatch(record_batch)
             }

arrow-ipc/src/reader/stream.rs

@@ -25,7 +25,7 @@ use arrow_data::UnsafeFlag;
 use arrow_schema::{ArrowError, SchemaRef};
 
 use crate::convert::MessageBuffer;
-use crate::reader::{RecordBatchDecoder, read_dictionary_impl};
+use crate::reader::{RecordBatchDecoder, message_custom_metadata, read_dictionary_impl};
 use crate::{CONTINUATION_MARKER, MessageHeader};
 
 /// A low-level interface for reading [`RecordBatch`] data from a stream of bytes
@@ -236,6 +236,7 @@ impl StreamDecoder {
                         }
                         MessageHeader::RecordBatch => {
                             let batch = message.header_as_record_batch().unwrap();
+                            let custom_metadata = message_custom_metadata(&message);
                             let schema = self.schema.clone().ok_or_else(|| {
                                 ArrowError::IpcError("Missing schema".to_string())
                             })?;
@@ -247,6 +248,7 @@ impl StreamDecoder {
                                 &version,
                             )?
                             .with_require_alignment(self.require_alignment)
+                            .with_custom_metadata(custom_metadata)
                             .read_record_batch()?;
                             self.state = DecoderState::default();
                             return Ok(Some(batch));