From 9390ec4fff1e16142911d34bb8816bdd2938522d Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Wed, 25 Feb 2026 12:51:28 -0500 Subject: [PATCH 1/3] feat: add `custom_metadata` support to `RecordBatch` with IPC read/write --- arrow-array/src/record_batch.rs | 129 +++++++++- arrow-flight/src/utils.rs | 9 + arrow-ipc/src/reader.rs | 46 +++- arrow-ipc/src/reader/stream.rs | 4 +- arrow-ipc/src/writer.rs | 227 ++++++++++++++++++ .../test_data/custom_metadata.arrow_file | Bin 0 -> 1130 bytes .../test_data/custom_metadata.arrow_stream | Bin 0 -> 880 bytes arrow-select/src/filter.rs | 10 +- arrow-select/src/take.rs | 2 + 9 files changed, 418 insertions(+), 9 deletions(-) create mode 100644 arrow-ipc/test_data/custom_metadata.arrow_file create mode 100644 arrow-ipc/test_data/custom_metadata.arrow_stream diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs index e05450a97f7f..3a8d260f1c4a 100644 --- a/arrow-array/src/record_batch.rs +++ b/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, } 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, 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, usize, HashMap) { + ( + 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 { + &self.custom_metadata + } + + /// Returns a mutable reference to the per-batch custom metadata. + pub fn custom_metadata_mut(&mut self) -> &mut HashMap { + &mut self.custom_metadata + } + + /// Sets the per-batch custom metadata, returning `self`. + pub fn with_custom_metadata(mut self, metadata: HashMap) -> Self { + self.custom_metadata = metadata; + self + } + /// Projects the schema onto the specified columns pub fn project(&self, indices: &[usize]) -> Result { 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 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); + } } diff --git a/arrow-flight/src/utils.rs b/arrow-flight/src/utils.rs index 6effb5f86aaf..604579988398 100644 --- a/arrow-flight/src/utils.rs +++ b/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) + } + }) })? } diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index aa66696271eb..f2f20ec2bbcf 100644 --- a/arrow-ipc/src/reader.rs +++ b/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 { + 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 /// 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, } 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) -> Self { + self.custom_metadata = custom_metadata; + self + } + /// Read the record batch, consuming the reader fn read_record_batch(mut self) -> Result { let mut variadic_counts: VecDeque = 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 { @@ -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 StreamReader { 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 StreamReader { .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) } diff --git a/arrow-ipc/src/reader/stream.rs b/arrow-ipc/src/reader/stream.rs index bfecf7b6ffa5..642d06950e96 100644 --- a/arrow-ipc/src/reader/stream.rs +++ b/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)); diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs index a05072a2c47c..10611028d02c 100644 --- a/arrow-ipc/src/writer.rs +++ b/arrow-ipc/src/writer.rs @@ -591,12 +591,19 @@ impl IpcDataGenerator { let b = batch_builder.finish(); b.as_union_value() }; + // serialize custom_metadata (must be created before MessageBuilder) + let fb_custom_metadata = (!batch.custom_metadata().is_empty()) + .then(|| crate::convert::metadata_to_fb(&mut fbb, batch.custom_metadata())); + // create an crate::Message let mut message = crate::MessageBuilder::new(&mut fbb); message.add_version(write_options.metadata_version); message.add_header_type(crate::MessageHeader::RecordBatch); message.add_bodyLength(arrow_data.len() as i64); message.add_header(root); + if let Some(m) = fb_custom_metadata { + message.add_custom_metadata(m); + } let root = message.finish(); fbb.finish(root, None); let finished_data = fbb.finished_data(); @@ -4320,4 +4327,224 @@ mod tests { let read_batch = reader.next().unwrap().unwrap(); assert_eq!(read_batch, batch2); } + + #[test] + fn test_custom_metadata_ipc_file_roundtrip() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata( + [("key".to_string(), "value".to_string())] + .into_iter() + .collect(), + ); + + // Write IPC file + let mut writer = FileWriter::try_new(vec![], &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + let data = writer.into_inner().unwrap(); + + // Read back + let reader = FileReader::try_new(Cursor::new(data), None).unwrap(); + let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); + assert_eq!(read_batches.len(), 1); + assert_eq!( + read_batches[0].custom_metadata().get("key"), + Some(&"value".to_string()) + ); + assert_eq!(read_batches[0], batch); + } + + #[test] + fn test_custom_metadata_ipc_stream_roundtrip() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata( + [("key".to_string(), "value".to_string())] + .into_iter() + .collect(), + ); + + // Write IPC stream + let mut buf = vec![]; + { + let mut writer = StreamWriter::try_new(&mut buf, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // Read back via StreamReader + let mut reader = StreamReader::try_new(Cursor::new(&buf), None).unwrap(); + let read_batch = reader.next().unwrap().unwrap(); + assert_eq!( + read_batch.custom_metadata().get("key"), + Some(&"value".to_string()) + ); + assert_eq!(read_batch, batch); + } + + #[test] + fn test_custom_metadata_stream_decoder_roundtrip() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata( + [("key".to_string(), "value".to_string())] + .into_iter() + .collect(), + ); + + // Write IPC stream + let mut buf = vec![]; + { + let mut writer = StreamWriter::try_new(&mut buf, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // Read back via StreamDecoder + let mut decoder = StreamDecoder::new(); + let mut buf = arrow_buffer::Buffer::from(buf); + let decoded = decoder.decode(&mut buf).unwrap().unwrap(); + assert_eq!( + decoded.custom_metadata().get("key"), + Some(&"value".to_string()) + ); + assert_eq!(decoded, batch); + } + + #[test] + fn test_multiple_batches_different_metadata() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata( + [("batch_index".to_string(), "0".to_string())] + .into_iter() + .collect(), + ); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![4, 5])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata( + [("batch_index".to_string(), "1".to_string())] + .into_iter() + .collect(), + ); + + // Write IPC file with both batches + let mut writer = FileWriter::try_new(vec![], &schema).unwrap(); + writer.write(&batch0).unwrap(); + writer.write(&batch1).unwrap(); + writer.finish().unwrap(); + let data = writer.into_inner().unwrap(); + + // Read back + let reader = FileReader::try_new(Cursor::new(data), None).unwrap(); + let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); + assert_eq!(read_batches.len(), 2); + assert_eq!( + read_batches[0].custom_metadata().get("batch_index"), + Some(&"0".to_string()) + ); + assert_eq!( + read_batches[1].custom_metadata().get("batch_index"), + Some(&"1".to_string()) + ); + } + + #[test] + fn test_no_custom_metadata_backward_compat() { + // Batches without custom metadata should still work fine + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap(); + assert!(batch.custom_metadata().is_empty()); + + let data = serialize_file(&batch); + let reader = FileReader::try_new(Cursor::new(data), None).unwrap(); + let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); + assert_eq!(read_batches.len(), 1); + assert!(read_batches[0].custom_metadata().is_empty()); + assert_eq!(read_batches[0], batch); + } + + #[test] + fn test_pyarrow_custom_metadata_file() { + let data = include_bytes!("../test_data/custom_metadata.arrow_file"); + let reader = FileReader::try_new(Cursor::new(data.as_slice()), None).unwrap(); + let batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); + assert_eq!(batches.len(), 2); + + // Batch 0 + assert_eq!( + batches[0].custom_metadata().get("batch_index"), + Some(&"0".to_string()) + ); + assert_eq!( + batches[0].custom_metadata().get("source"), + Some(&"test".to_string()) + ); + assert_eq!(batches[0].num_rows(), 3); + + // Batch 1 + assert_eq!( + batches[1].custom_metadata().get("batch_index"), + Some(&"1".to_string()) + ); + assert_eq!( + batches[1].custom_metadata().get("source"), + Some(&"test".to_string()) + ); + assert_eq!(batches[1].num_rows(), 2); + } + + #[test] + fn test_pyarrow_custom_metadata_stream() { + let data = include_bytes!("../test_data/custom_metadata.arrow_stream"); + let mut reader = StreamReader::try_new(Cursor::new(data.as_slice()), None).unwrap(); + + let batch0 = reader.next().unwrap().unwrap(); + assert_eq!( + batch0.custom_metadata().get("batch_index"), + Some(&"0".to_string()) + ); + assert_eq!( + batch0.custom_metadata().get("source"), + Some(&"test".to_string()) + ); + assert_eq!(batch0.num_rows(), 3); + + let batch1 = reader.next().unwrap().unwrap(); + assert_eq!( + batch1.custom_metadata().get("batch_index"), + Some(&"1".to_string()) + ); + assert_eq!( + batch1.custom_metadata().get("source"), + Some(&"test".to_string()) + ); + assert_eq!(batch1.num_rows(), 2); + + assert!(reader.next().is_none()); + } } diff --git a/arrow-ipc/test_data/custom_metadata.arrow_file b/arrow-ipc/test_data/custom_metadata.arrow_file new file mode 100644 index 0000000000000000000000000000000000000000..e7982d807ad6665ee6e04ec9a207408943ab90d6 GIT binary patch literal 1130 zcmd5+y-LGS6#mjCX{n)P=+MC-qhkdJ2bY4QAarq*GznPH3O0g+59A~GAdWtOZ=v6J z?>!iTh=PL$&OPTlIe+)$B+J$6{NfOJCp`hk$-o2|QY08tBSkSWh{inQZ~_COQ)dBQ z7*EGU@%zC>L2qDN-ZWm19nLXhjua!8H4inTK|facH?`_!^XndmtQ0d5_5I?Ui|EK^JbDT}#1C!0z_`u43n^uUqT($|XOR#)7;@ zUYBikb-CHr%>&VdJ$Zl4?wYKaA=xM~SDR{vHI@&b#3Lk zQrNHOTev8H{MCH9mOq}aY8rhvU%f+-Nsm6U&ENK_n^$@lE4Gid8&==uD!-b0B+}kW w>OE%s4pO&W+E`JHGoB|-^SPg8CLJpo$@X|23Y)!fPyO+qzSB3hv~>TN51VF5c>n+a literal 0 HcmV?d00001 diff --git a/arrow-ipc/test_data/custom_metadata.arrow_stream b/arrow-ipc/test_data/custom_metadata.arrow_stream new file mode 100644 index 0000000000000000000000000000000000000000..b6b40a6c39e098e715478521893db3eb313f1dcf GIT binary patch literal 880 zcmb_aL2AQ53>+tR64DTS$f1WGe9SS0JfZ&xj%{e5O`&n1hd!te>al;>%S^-Re)SI}W&^rn+rPc5lvjR=RXfCb8`j_Us=Su(NTPQudjVd$I=}z` literal 0 HcmV?d00001 diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index e95d01f2b592..f84926df33a2 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -395,12 +395,12 @@ impl FilterPredicate { // SAFETY: we know that the set of filtered arrays will match the schema of the original // record batch + let custom_metadata = record_batch.custom_metadata().clone(); unsafe { - Ok(RecordBatch::new_unchecked( - record_batch.schema(), - filtered_arrays, - self.count, - )) + Ok( + RecordBatch::new_unchecked(record_batch.schema(), filtered_arrays, self.count) + .with_custom_metadata(custom_metadata), + ) } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index cbb65ac915dd..8bc948dc4f32 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1116,7 +1116,9 @@ pub fn take_record_batch( .iter() .map(|c| take(c, indices, None)) .collect::, _>>()?; + let custom_metadata = record_batch.custom_metadata().clone(); RecordBatch::try_new(record_batch.schema(), columns) + .map(|b| b.with_custom_metadata(custom_metadata)) } #[cfg(test)] From 49009dbaf6d6011ec9bae970f6bb29aaa00c5ac4 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 11 May 2026 12:24:33 -0400 Subject: [PATCH 2/3] refactor: store RecordBatch custom_metadata as Option> Addresses @jhorstmann's review: the previous inline HashMap pushed RecordBatch from 40 to 88 bytes and made every clone copy the map. Using Option>> brings the struct to 48 bytes (one niche-optimized pointer) and clones become refcount bumps. API surface reduced to four methods on RecordBatch: - custom_metadata() -> Option<&HashMap> - custom_metadata_mut() -> &mut HashMap - with_custom_metadata(HashMap) -> Self - into_parts_with_custom_metadata() -> (..., Option>) with_custom_metadata normalizes an empty map to None. Because custom_metadata_mut cannot observe a subsequent .clear(), PartialEq is now implemented by hand to treat Some(empty) and None as equal, keeping the two paths to "no metadata" indistinguishable to callers. message_custom_metadata in arrow-ipc now returns Option directly; downstream callers in arrow-ipc, arrow-select and arrow-flight clone the HashMap on transfer (typically tiny). Co-Authored-By: Claude Opus 4.7 (1M context) --- arrow-array/src/record_batch.rs | 140 +++++++++++++++++++++++++------- arrow-flight/src/utils.rs | 9 +- arrow-ipc/src/reader.rs | 41 +++++----- arrow-ipc/src/writer.rs | 70 +++++++--------- arrow-select/src/filter.rs | 12 +-- arrow-select/src/take.rs | 8 +- 6 files changed, 176 insertions(+), 104 deletions(-) diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs index 3a8d260f1c4a..c321f17e3637 100644 --- a/arrow-array/src/record_batch.rs +++ b/arrow-array/src/record_batch.rs @@ -221,7 +221,7 @@ macro_rules! record_batch { /// ("c", Utf8, ["alpha", "beta", "gamma"]) /// ); /// ``` -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug)] pub struct RecordBatch { schema: SchemaRef, columns: Vec>, @@ -234,8 +234,34 @@ pub struct RecordBatch { /// 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, + /// flatbuffer, allowing per-batch metadata separate from schema-level + /// metadata. Stored as `Option>` so that a `RecordBatch` without + /// custom metadata adds only a pointer's worth of overhead, and clones + /// share the map. + custom_metadata: Option>>, +} + +// Custom equality: a batch built with an empty metadata map compares equal to +// one built without any metadata at all. This is also reachable via +// `custom_metadata_mut().clear()`, so we collapse the two forms here rather +// than letting the derived impl surface the distinction. +impl PartialEq for RecordBatch { + fn eq(&self, other: &Self) -> bool { + if self.row_count != other.row_count || self.schema != other.schema { + return false; + } + if self.columns != other.columns { + return false; + } + match ( + self.custom_metadata.as_deref(), + other.custom_metadata.as_deref(), + ) { + (None, None) => true, + (Some(m), None) | (None, Some(m)) => m.is_empty(), + (Some(a), Some(b)) => a == b, + } + } } impl RecordBatch { @@ -296,7 +322,7 @@ impl RecordBatch { schema, columns, row_count, - custom_metadata: HashMap::new(), + custom_metadata: None, } } @@ -324,7 +350,7 @@ impl RecordBatch { schema, columns, row_count: 0, - custom_metadata: HashMap::new(), + custom_metadata: None, } } @@ -399,7 +425,7 @@ impl RecordBatch { schema, columns, row_count, - custom_metadata: HashMap::new(), + custom_metadata: None, }) } @@ -411,15 +437,25 @@ impl RecordBatch { (self.schema, self.columns, self.row_count) } - /// Return the schema, columns, row count and custom metadata of this [`RecordBatch`] + /// Return the schema, columns, row count and custom metadata of this [`RecordBatch`]. + /// + /// The returned metadata is `None` if this batch has no custom metadata. + /// If the batch uniquely owns its metadata (the common case), the map is + /// returned without cloning; otherwise it is cloned out of the shared + /// reference. pub fn into_parts_with_custom_metadata( self, - ) -> (SchemaRef, Vec, usize, HashMap) { + ) -> ( + SchemaRef, + Vec, + usize, + Option>, + ) { ( self.schema, self.columns, self.row_count, - self.custom_metadata, + self.custom_metadata.map(Arc::unwrap_or_clone), ) } @@ -477,22 +513,39 @@ impl RecordBatch { &mut schema.metadata } - /// Returns a reference to the per-batch custom metadata. + /// Returns the per-batch custom metadata, or `None` if not set. /// - /// This metadata corresponds to the `custom_metadata` field on the IPC - /// `Message` flatbuffer, separate from schema-level metadata. - pub fn custom_metadata(&self) -> &HashMap { - &self.custom_metadata + /// This corresponds to the `custom_metadata` field on the IPC `Message` + /// flatbuffer, separate from schema-level metadata. + pub fn custom_metadata(&self) -> Option<&HashMap> { + self.custom_metadata.as_deref() } - /// Returns a mutable reference to the per-batch custom metadata. + /// Returns a mutable reference to the per-batch custom metadata, allocating + /// an empty map on first access. + /// + /// Cheap if this [`RecordBatch`] uniquely owns the metadata; otherwise the + /// underlying map is cloned via [`Arc::make_mut`]. An empty map left after + /// clearing entries still reports as `Some(_)` from + /// [`Self::custom_metadata`]; two batches that differ only in this respect + /// still compare equal. pub fn custom_metadata_mut(&mut self) -> &mut HashMap { - &mut self.custom_metadata + Arc::make_mut( + self.custom_metadata + .get_or_insert_with(|| Arc::new(HashMap::new())), + ) } /// Sets the per-batch custom metadata, returning `self`. + /// + /// An empty map is normalized to "no metadata", so a batch built with an + /// empty map compares equal to one built without calling this method. pub fn with_custom_metadata(mut self, metadata: HashMap) -> Self { - self.custom_metadata = metadata; + self.custom_metadata = if metadata.is_empty() { + None + } else { + Some(Arc::new(metadata)) + }; self } @@ -516,12 +569,13 @@ impl RecordBatch { // Since we're starting from a valid RecordBatch and project // creates a strict subset of the original, there's no need to // redo the validation checks in `try_new_with_options`. - Ok(RecordBatch::new_unchecked( + let mut projected = RecordBatch::new_unchecked( SchemaRef::new(projected_schema), batch_fields, self.row_count, - ) - .with_custom_metadata(self.custom_metadata.clone())) + ); + projected.custom_metadata = self.custom_metadata.clone(); + Ok(projected) } } @@ -617,8 +671,10 @@ 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)) + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).map(|mut b| { + b.custom_metadata = custom_metadata; + b + }) } /// Returns the number of columns in the record batch. @@ -913,7 +969,7 @@ impl From for RecordBatch { schema: Arc::new(Schema::new(fields)), row_count, columns, - custom_metadata: HashMap::new(), + custom_metadata: None, } } } @@ -1846,12 +1902,12 @@ mod tests { #[test] fn test_with_custom_metadata() { let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); - assert!(batch.custom_metadata().is_empty()); + assert!(batch.custom_metadata().is_none()); 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); + assert_eq!(batch.custom_metadata(), Some(&metadata)); } #[test] @@ -1861,7 +1917,7 @@ mod tests { .custom_metadata_mut() .insert("key".to_string(), "value".to_string()); assert_eq!( - batch.custom_metadata().get("key"), + batch.custom_metadata().and_then(|m| m.get("key")), Some(&"value".to_string()) ); } @@ -1874,7 +1930,7 @@ mod tests { let batch = batch.with_custom_metadata(metadata.clone()); let sliced = batch.slice(0, 2); - assert_eq!(sliced.custom_metadata(), &metadata); + assert_eq!(sliced.custom_metadata(), Some(&metadata)); } #[test] @@ -1888,7 +1944,7 @@ mod tests { let batch = batch.with_custom_metadata(metadata.clone()); let projected = batch.project(&[0]).unwrap(); - assert_eq!(projected.custom_metadata(), &metadata); + assert_eq!(projected.custom_metadata(), Some(&metadata)); } #[test] @@ -1902,7 +1958,7 @@ mod tests { assert_eq!(schema.fields().len(), 1); assert_eq!(columns.len(), 1); assert_eq!(row_count, 3); - assert_eq!(custom_metadata, metadata); + assert_eq!(custom_metadata, Some(metadata)); } #[test] @@ -1919,4 +1975,30 @@ mod tests { let batch1 = batch1.with_custom_metadata(metadata); assert_ne!(batch1, batch2); } + + #[test] + fn test_empty_custom_metadata_normalized_to_none() { + // A batch built with an empty map compares equal to one with no + // metadata, and the setter normalizes the empty map away. + let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); + let batch2 = record_batch!(("a", Int32, [1, 2, 3])) + .unwrap() + .with_custom_metadata(HashMap::new()); + assert_eq!(batch1, batch2); + assert!(batch2.custom_metadata().is_none()); + } + + #[test] + fn test_equality_after_mut_clear() { + // `custom_metadata_mut().clear()` cannot normalize the Option back to + // None, so it leaves an empty `Some` in place. PartialEq must still + // treat this as equal to a batch with no metadata. + let mut metadata = HashMap::new(); + metadata.insert("k".to_string(), "v".to_string()); + let no_meta = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); + let mut cleared = no_meta.clone().with_custom_metadata(metadata); + cleared.custom_metadata_mut().clear(); + assert!(cleared.custom_metadata().is_some_and(|m| m.is_empty())); + assert_eq!(no_meta, cleared); + } } diff --git a/arrow-flight/src/utils.rs b/arrow-flight/src/utils.rs index 604579988398..94042c0a13f2 100644 --- a/arrow-flight/src/utils.rs +++ b/arrow-flight/src/utils.rs @@ -79,12 +79,9 @@ 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) - } + .map(|rb| match custom_metadata { + Some(m) => rb.with_custom_metadata(m), + None => rb, }) })? } diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index d0c6fc079101..0667181b6162 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -49,17 +49,20 @@ 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 { - 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()); - } +/// Returns `None` if the message has no custom metadata. +pub fn message_custom_metadata(message: &crate::Message) -> Option> { + let list = message.custom_metadata()?; + let mut metadata = HashMap::with_capacity(list.len()); + for kv in list { + if let (Some(k), Some(v)) = (kv.key(), kv.value()) { + metadata.insert(k.to_string(), v.to_string()); } } - metadata + if metadata.is_empty() { + None + } else { + Some(metadata) + } } /// Read a buffer based on offset and length @@ -486,7 +489,7 @@ 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, + custom_metadata: Option>, } impl<'a> RecordBatchDecoder<'a> { @@ -523,7 +526,7 @@ impl<'a> RecordBatchDecoder<'a> { projection: None, require_alignment: false, skip_validation: UnsafeFlag::new(), - custom_metadata: HashMap::new(), + custom_metadata: None, }) } @@ -563,7 +566,10 @@ impl<'a> RecordBatchDecoder<'a> { } /// Set per-batch custom metadata to attach to the decoded [`RecordBatch`] - pub(crate) fn with_custom_metadata(mut self, custom_metadata: HashMap) -> Self { + pub(crate) fn with_custom_metadata( + mut self, + custom_metadata: Option>, + ) -> Self { self.custom_metadata = custom_metadata; self } @@ -578,7 +584,7 @@ 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 custom_metadata = self.custom_metadata.take(); let schema = Arc::clone(&self.schema); let batch = if let Some(projection) = self.projection { @@ -635,12 +641,9 @@ impl<'a> RecordBatchDecoder<'a> { } }; - batch.map(|b| { - if custom_metadata.is_empty() { - b - } else { - b.with_custom_metadata(custom_metadata) - } + batch.map(|b| match custom_metadata { + Some(m) => b.with_custom_metadata(m), + None => b, }) } diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs index 1c9c47c50937..4db1d9ee6f4d 100644 --- a/arrow-ipc/src/writer.rs +++ b/arrow-ipc/src/writer.rs @@ -592,8 +592,10 @@ impl IpcDataGenerator { b.as_union_value() }; // serialize custom_metadata (must be created before MessageBuilder) - let fb_custom_metadata = (!batch.custom_metadata().is_empty()) - .then(|| crate::convert::metadata_to_fb(&mut fbb, batch.custom_metadata())); + let fb_custom_metadata = batch + .custom_metadata() + .filter(|m| !m.is_empty()) + .map(|m| crate::convert::metadata_to_fb(&mut fbb, m)); // create an crate::Message let mut message = crate::MessageBuilder::new(&mut fbb); @@ -4328,6 +4330,10 @@ mod tests { assert_eq!(read_batch, batch2); } + fn map_of>(items: I) -> HashMap { + items.into_iter().collect() + } + #[test] fn test_custom_metadata_ipc_file_roundtrip() { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -4336,11 +4342,7 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], ) .unwrap() - .with_custom_metadata( - [("key".to_string(), "value".to_string())] - .into_iter() - .collect(), - ); + .with_custom_metadata(map_of([("key".to_string(), "value".to_string())])); // Write IPC file let mut writer = FileWriter::try_new(vec![], &schema).unwrap(); @@ -4353,7 +4355,7 @@ mod tests { let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); assert_eq!(read_batches.len(), 1); assert_eq!( - read_batches[0].custom_metadata().get("key"), + read_batches[0].custom_metadata().and_then(|m| m.get("key")), Some(&"value".to_string()) ); assert_eq!(read_batches[0], batch); @@ -4367,11 +4369,7 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], ) .unwrap() - .with_custom_metadata( - [("key".to_string(), "value".to_string())] - .into_iter() - .collect(), - ); + .with_custom_metadata(map_of([("key".to_string(), "value".to_string())])); // Write IPC stream let mut buf = vec![]; @@ -4385,7 +4383,7 @@ mod tests { let mut reader = StreamReader::try_new(Cursor::new(&buf), None).unwrap(); let read_batch = reader.next().unwrap().unwrap(); assert_eq!( - read_batch.custom_metadata().get("key"), + read_batch.custom_metadata().and_then(|m| m.get("key")), Some(&"value".to_string()) ); assert_eq!(read_batch, batch); @@ -4399,11 +4397,7 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], ) .unwrap() - .with_custom_metadata( - [("key".to_string(), "value".to_string())] - .into_iter() - .collect(), - ); + .with_custom_metadata(map_of([("key".to_string(), "value".to_string())])); // Write IPC stream let mut buf = vec![]; @@ -4418,7 +4412,7 @@ mod tests { let mut buf = arrow_buffer::Buffer::from(buf); let decoded = decoder.decode(&mut buf).unwrap().unwrap(); assert_eq!( - decoded.custom_metadata().get("key"), + decoded.custom_metadata().and_then(|m| m.get("key")), Some(&"value".to_string()) ); assert_eq!(decoded, batch); @@ -4432,21 +4426,13 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], ) .unwrap() - .with_custom_metadata( - [("batch_index".to_string(), "0".to_string())] - .into_iter() - .collect(), - ); + .with_custom_metadata(map_of([("batch_index".to_string(), "0".to_string())])); let batch1 = RecordBatch::try_new( schema.clone(), vec![Arc::new(Int32Array::from(vec![4, 5])) as ArrayRef], ) .unwrap() - .with_custom_metadata( - [("batch_index".to_string(), "1".to_string())] - .into_iter() - .collect(), - ); + .with_custom_metadata(map_of([("batch_index".to_string(), "1".to_string())])); // Write IPC file with both batches let mut writer = FileWriter::try_new(vec![], &schema).unwrap(); @@ -4460,11 +4446,11 @@ mod tests { let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); assert_eq!(read_batches.len(), 2); assert_eq!( - read_batches[0].custom_metadata().get("batch_index"), + read_batches[0].custom_metadata().and_then(|m| m.get("batch_index")), Some(&"0".to_string()) ); assert_eq!( - read_batches[1].custom_metadata().get("batch_index"), + read_batches[1].custom_metadata().and_then(|m| m.get("batch_index")), Some(&"1".to_string()) ); } @@ -4478,13 +4464,13 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], ) .unwrap(); - assert!(batch.custom_metadata().is_empty()); + assert!(batch.custom_metadata().is_none()); let data = serialize_file(&batch); let reader = FileReader::try_new(Cursor::new(data), None).unwrap(); let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); assert_eq!(read_batches.len(), 1); - assert!(read_batches[0].custom_metadata().is_empty()); + assert!(read_batches[0].custom_metadata().is_none()); assert_eq!(read_batches[0], batch); } @@ -4497,22 +4483,22 @@ mod tests { // Batch 0 assert_eq!( - batches[0].custom_metadata().get("batch_index"), + batches[0].custom_metadata().and_then(|m| m.get("batch_index")), Some(&"0".to_string()) ); assert_eq!( - batches[0].custom_metadata().get("source"), + batches[0].custom_metadata().and_then(|m| m.get("source")), Some(&"test".to_string()) ); assert_eq!(batches[0].num_rows(), 3); // Batch 1 assert_eq!( - batches[1].custom_metadata().get("batch_index"), + batches[1].custom_metadata().and_then(|m| m.get("batch_index")), Some(&"1".to_string()) ); assert_eq!( - batches[1].custom_metadata().get("source"), + batches[1].custom_metadata().and_then(|m| m.get("source")), Some(&"test".to_string()) ); assert_eq!(batches[1].num_rows(), 2); @@ -4525,22 +4511,22 @@ mod tests { let batch0 = reader.next().unwrap().unwrap(); assert_eq!( - batch0.custom_metadata().get("batch_index"), + batch0.custom_metadata().and_then(|m| m.get("batch_index")), Some(&"0".to_string()) ); assert_eq!( - batch0.custom_metadata().get("source"), + batch0.custom_metadata().and_then(|m| m.get("source")), Some(&"test".to_string()) ); assert_eq!(batch0.num_rows(), 3); let batch1 = reader.next().unwrap().unwrap(); assert_eq!( - batch1.custom_metadata().get("batch_index"), + batch1.custom_metadata().and_then(|m| m.get("batch_index")), Some(&"1".to_string()) ); assert_eq!( - batch1.custom_metadata().get("source"), + batch1.custom_metadata().and_then(|m| m.get("source")), Some(&"test".to_string()) ); assert_eq!(batch1.num_rows(), 2); diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index f84926df33a2..0d614ba584c6 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -395,12 +395,14 @@ impl FilterPredicate { // SAFETY: we know that the set of filtered arrays will match the schema of the original // record batch - let custom_metadata = record_batch.custom_metadata().clone(); + let custom_metadata = record_batch.custom_metadata().cloned(); unsafe { - Ok( - RecordBatch::new_unchecked(record_batch.schema(), filtered_arrays, self.count) - .with_custom_metadata(custom_metadata), - ) + let mut batch = + RecordBatch::new_unchecked(record_batch.schema(), filtered_arrays, self.count); + if let Some(m) = custom_metadata { + batch = batch.with_custom_metadata(m); + } + Ok(batch) } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 8bc948dc4f32..3dd32894fae2 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1116,9 +1116,11 @@ pub fn take_record_batch( .iter() .map(|c| take(c, indices, None)) .collect::, _>>()?; - let custom_metadata = record_batch.custom_metadata().clone(); - RecordBatch::try_new(record_batch.schema(), columns) - .map(|b| b.with_custom_metadata(custom_metadata)) + let custom_metadata = record_batch.custom_metadata().cloned(); + RecordBatch::try_new(record_batch.schema(), columns).map(|b| match custom_metadata { + Some(m) => b.with_custom_metadata(m), + None => b, + }) } #[cfg(test)] From c3e53303e6149b293820a37837d7977ea1291019 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 11 May 2026 15:49:18 -0400 Subject: [PATCH 3/3] refactor: introduce CustomMetadata alias and tighten RecordBatch API Address review feedback from @viirya on #9445: - `RecordBatch::with_custom_metadata` now accepts `Arc>` so the same metadata map can be shared across many batches without cloning. - `into_parts_with_custom_metadata` returns the shared `Arc` instead of unwrap-cloning out of it. - `PartialEq` now compares custom_metadata before columns: in the common case where neither batch has metadata it is a near-free `None == None` check that short-circuits real mismatches without scanning arrays. - Introduce `pub type CustomMetadata = Arc>` and use it on the field, setter, and parts accessor. Avoids the type-complexity clippy lint and gives the type a name callers can refer to. - Update IPC reader, arrow-select filter/take, and arrow-flight callsites to wrap their cloned maps in `Arc::new` when reattaching. Co-Authored-By: Claude Opus 4.7 (1M context) --- arrow-array/src/lib.rs | 3 +- arrow-array/src/record_batch.rs | 59 ++++++++++++++++++--------------- arrow-flight/src/utils.rs | 2 +- arrow-ipc/src/reader.rs | 2 +- arrow-ipc/src/writer.rs | 20 +++++++---- arrow-select/src/filter.rs | 2 +- arrow-select/src/take.rs | 2 +- 7 files changed, 53 insertions(+), 37 deletions(-) diff --git a/arrow-array/src/lib.rs b/arrow-array/src/lib.rs index a5f9bf5e711c..3424a33f2ffe 100644 --- a/arrow-array/src/lib.rs +++ b/arrow-array/src/lib.rs @@ -234,7 +234,8 @@ pub use array::*; mod record_batch; pub use record_batch::{ - RecordBatch, RecordBatchIterator, RecordBatchOptions, RecordBatchReader, RecordBatchWriter, + CustomMetadata, RecordBatch, RecordBatchIterator, RecordBatchOptions, RecordBatchReader, + RecordBatchWriter, }; mod arithmetic; diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs index c321f17e3637..5b320d3b1368 100644 --- a/arrow-array/src/record_batch.rs +++ b/arrow-array/src/record_batch.rs @@ -36,6 +36,14 @@ pub trait RecordBatchReader: Iterator> { fn schema(&self) -> SchemaRef; } +/// Shared per-batch custom metadata. +/// +/// This is the type stored inside a [`RecordBatch`] and accepted by +/// [`RecordBatch::with_custom_metadata`]. Cloning is cheap (an [`Arc`] +/// refcount bump), so the same metadata map can be attached to many batches +/// without copying the underlying [`HashMap`]. +pub type CustomMetadata = Arc>; + impl RecordBatchReader for Box { fn schema(&self) -> SchemaRef { self.as_ref().schema() @@ -238,7 +246,7 @@ pub struct RecordBatch { /// metadata. Stored as `Option>` so that a `RecordBatch` without /// custom metadata adds only a pointer's worth of overhead, and clones /// share the map. - custom_metadata: Option>>, + custom_metadata: Option, } // Custom equality: a batch built with an empty metadata map compares equal to @@ -250,17 +258,18 @@ impl PartialEq for RecordBatch { if self.row_count != other.row_count || self.schema != other.schema { return false; } - if self.columns != other.columns { - return false; - } - match ( + let metadata_eq = match ( self.custom_metadata.as_deref(), other.custom_metadata.as_deref(), ) { (None, None) => true, (Some(m), None) | (None, Some(m)) => m.is_empty(), (Some(a), Some(b)) => a == b, + }; + if !metadata_eq { + return false; } + self.columns == other.columns } } @@ -439,23 +448,17 @@ impl RecordBatch { /// Return the schema, columns, row count and custom metadata of this [`RecordBatch`]. /// - /// The returned metadata is `None` if this batch has no custom metadata. - /// If the batch uniquely owns its metadata (the common case), the map is - /// returned without cloning; otherwise it is cloned out of the shared - /// reference. + /// The returned metadata is `None` if this batch has no custom metadata, + /// otherwise it is the shared [`Arc`] held by this batch. Callers that need + /// an owned `HashMap` can use [`Arc::unwrap_or_clone`]. pub fn into_parts_with_custom_metadata( self, - ) -> ( - SchemaRef, - Vec, - usize, - Option>, - ) { + ) -> (SchemaRef, Vec, usize, Option) { ( self.schema, self.columns, self.row_count, - self.custom_metadata.map(Arc::unwrap_or_clone), + self.custom_metadata, ) } @@ -538,13 +541,17 @@ impl RecordBatch { /// Sets the per-batch custom metadata, returning `self`. /// + /// Takes an [`Arc`] so the same metadata map can be shared across many + /// [`RecordBatch`]es without cloning. Callers with a fresh `HashMap` + /// can wrap it via [`Arc::new`]. + /// /// An empty map is normalized to "no metadata", so a batch built with an /// empty map compares equal to one built without calling this method. - pub fn with_custom_metadata(mut self, metadata: HashMap) -> Self { + pub fn with_custom_metadata(mut self, metadata: CustomMetadata) -> Self { self.custom_metadata = if metadata.is_empty() { None } else { - Some(Arc::new(metadata)) + Some(metadata) }; self } @@ -1906,7 +1913,7 @@ mod tests { let mut metadata = HashMap::new(); metadata.insert("key".to_string(), "value".to_string()); - let batch = batch.with_custom_metadata(metadata.clone()); + let batch = batch.with_custom_metadata(Arc::new(metadata.clone())); assert_eq!(batch.custom_metadata(), Some(&metadata)); } @@ -1927,7 +1934,7 @@ mod tests { 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 batch = batch.with_custom_metadata(Arc::new(metadata.clone())); let sliced = batch.slice(0, 2); assert_eq!(sliced.custom_metadata(), Some(&metadata)); @@ -1941,7 +1948,7 @@ mod tests { let mut metadata = HashMap::new(); metadata.insert("k".to_string(), "v".to_string()); - let batch = batch.with_custom_metadata(metadata.clone()); + let batch = batch.with_custom_metadata(Arc::new(metadata.clone())); let projected = batch.project(&[0]).unwrap(); assert_eq!(projected.custom_metadata(), Some(&metadata)); @@ -1952,13 +1959,13 @@ mod tests { 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 batch = batch.with_custom_metadata(Arc::new(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, Some(metadata)); + assert_eq!(custom_metadata.as_deref(), Some(&metadata)); } #[test] @@ -1972,7 +1979,7 @@ mod tests { // Different metadata -> not equal let mut metadata = HashMap::new(); metadata.insert("k".to_string(), "v".to_string()); - let batch1 = batch1.with_custom_metadata(metadata); + let batch1 = batch1.with_custom_metadata(Arc::new(metadata)); assert_ne!(batch1, batch2); } @@ -1983,7 +1990,7 @@ mod tests { let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); let batch2 = record_batch!(("a", Int32, [1, 2, 3])) .unwrap() - .with_custom_metadata(HashMap::new()); + .with_custom_metadata(Arc::new(HashMap::new())); assert_eq!(batch1, batch2); assert!(batch2.custom_metadata().is_none()); } @@ -1996,7 +2003,7 @@ mod tests { let mut metadata = HashMap::new(); metadata.insert("k".to_string(), "v".to_string()); let no_meta = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); - let mut cleared = no_meta.clone().with_custom_metadata(metadata); + let mut cleared = no_meta.clone().with_custom_metadata(Arc::new(metadata)); cleared.custom_metadata_mut().clear(); assert!(cleared.custom_metadata().is_some_and(|m| m.is_empty())); assert_eq!(no_meta, cleared); diff --git a/arrow-flight/src/utils.rs b/arrow-flight/src/utils.rs index 94042c0a13f2..5714b78509e7 100644 --- a/arrow-flight/src/utils.rs +++ b/arrow-flight/src/utils.rs @@ -80,7 +80,7 @@ pub fn flight_data_to_arrow_batch( &message.version(), ) .map(|rb| match custom_metadata { - Some(m) => rb.with_custom_metadata(m), + Some(m) => rb.with_custom_metadata(Arc::new(m)), None => rb, }) })? diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index 0667181b6162..979167a5d130 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -642,7 +642,7 @@ impl<'a> RecordBatchDecoder<'a> { }; batch.map(|b| match custom_metadata { - Some(m) => b.with_custom_metadata(m), + Some(m) => b.with_custom_metadata(Arc::new(m)), None => b, }) } diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs index 4db1d9ee6f4d..fdb8138902d1 100644 --- a/arrow-ipc/src/writer.rs +++ b/arrow-ipc/src/writer.rs @@ -4330,8 +4330,8 @@ mod tests { assert_eq!(read_batch, batch2); } - fn map_of>(items: I) -> HashMap { - items.into_iter().collect() + fn map_of>(items: I) -> CustomMetadata { + Arc::new(items.into_iter().collect()) } #[test] @@ -4446,11 +4446,15 @@ mod tests { let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect(); assert_eq!(read_batches.len(), 2); assert_eq!( - read_batches[0].custom_metadata().and_then(|m| m.get("batch_index")), + read_batches[0] + .custom_metadata() + .and_then(|m| m.get("batch_index")), Some(&"0".to_string()) ); assert_eq!( - read_batches[1].custom_metadata().and_then(|m| m.get("batch_index")), + read_batches[1] + .custom_metadata() + .and_then(|m| m.get("batch_index")), Some(&"1".to_string()) ); } @@ -4483,7 +4487,9 @@ mod tests { // Batch 0 assert_eq!( - batches[0].custom_metadata().and_then(|m| m.get("batch_index")), + batches[0] + .custom_metadata() + .and_then(|m| m.get("batch_index")), Some(&"0".to_string()) ); assert_eq!( @@ -4494,7 +4500,9 @@ mod tests { // Batch 1 assert_eq!( - batches[1].custom_metadata().and_then(|m| m.get("batch_index")), + batches[1] + .custom_metadata() + .and_then(|m| m.get("batch_index")), Some(&"1".to_string()) ); assert_eq!( diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 0d614ba584c6..ad7c5fdfed81 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -400,7 +400,7 @@ impl FilterPredicate { let mut batch = RecordBatch::new_unchecked(record_batch.schema(), filtered_arrays, self.count); if let Some(m) = custom_metadata { - batch = batch.with_custom_metadata(m); + batch = batch.with_custom_metadata(Arc::new(m)); } Ok(batch) } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 3dd32894fae2..67fe850f738c 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1118,7 +1118,7 @@ pub fn take_record_batch( .collect::, _>>()?; let custom_metadata = record_batch.custom_metadata().cloned(); RecordBatch::try_new(record_batch.schema(), columns).map(|b| match custom_metadata { - Some(m) => b.with_custom_metadata(m), + Some(m) => b.with_custom_metadata(Arc::new(m)), None => b, }) }