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 e05450a97f7f..5b320d3b1368 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; @@ -35,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() @@ -220,7 +229,7 @@ macro_rules! record_batch { /// ("c", Utf8, ["alpha", "beta", "gamma"]) /// ); /// ``` -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug)] pub struct RecordBatch { schema: SchemaRef, columns: Vec>, @@ -229,6 +238,39 @@ 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. 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; + } + 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 + } } impl RecordBatch { @@ -289,6 +331,7 @@ impl RecordBatch { schema, columns, row_count, + custom_metadata: None, } } @@ -316,6 +359,7 @@ impl RecordBatch { schema, columns, row_count: 0, + custom_metadata: None, } } @@ -390,14 +434,34 @@ impl RecordBatch { schema, columns, row_count, + custom_metadata: None, }) } /// 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`]. + /// + /// 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) { + ( + 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 +480,7 @@ impl RecordBatch { schema, columns: self.columns, row_count: self.row_count, + custom_metadata: self.custom_metadata, }) } @@ -451,6 +516,46 @@ impl RecordBatch { &mut schema.metadata } + /// Returns the per-batch custom metadata, or `None` if not set. + /// + /// 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, 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 { + Arc::make_mut( + self.custom_metadata + .get_or_insert_with(|| Arc::new(HashMap::new())), + ) + } + + /// 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: CustomMetadata) -> Self { + self.custom_metadata = if metadata.is_empty() { + None + } else { + Some(metadata) + }; + self + } + /// Projects the schema onto the specified columns pub fn project(&self, indices: &[usize]) -> Result { let projected_schema = self.schema.project(indices)?; @@ -471,11 +576,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, - )) + ); + projected.custom_metadata = self.custom_metadata.clone(); + Ok(projected) } } @@ -570,7 +677,11 @@ impl RecordBatch { } } } - RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) + let custom_metadata = self.custom_metadata.clone(); + 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. @@ -691,6 +802,7 @@ impl RecordBatch { schema: self.schema.clone(), columns, row_count: length, + custom_metadata: self.custom_metadata.clone(), } } @@ -864,6 +976,7 @@ impl From for RecordBatch { schema: Arc::new(Schema::new(fields)), row_count, columns, + custom_metadata: None, } } } @@ -1792,4 +1905,107 @@ 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_none()); + + let mut metadata = HashMap::new(); + metadata.insert("key".to_string(), "value".to_string()); + let batch = batch.with_custom_metadata(Arc::new(metadata.clone())); + assert_eq!(batch.custom_metadata(), Some(&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().and_then(|m| m.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(Arc::new(metadata.clone())); + + let sliced = batch.slice(0, 2); + assert_eq!(sliced.custom_metadata(), Some(&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(Arc::new(metadata.clone())); + + let projected = batch.project(&[0]).unwrap(); + assert_eq!(projected.custom_metadata(), Some(&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(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.as_deref(), Some(&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(Arc::new(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(Arc::new(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(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 6effb5f86aaf..5714b78509e7 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,10 @@ pub fn flight_data_to_arrow_batch( None, &message.version(), ) + .map(|rb| match custom_metadata { + 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 1d5e06c6871c..979167a5d130 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -47,6 +47,24 @@ 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 `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()); + } + } + if metadata.is_empty() { + None + } else { + Some(metadata) + } +} + /// Read a buffer based on offset and length /// From /// Each constituent buffer is first compressed with the indicated @@ -470,6 +488,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: Option>, } impl<'a> RecordBatchDecoder<'a> { @@ -506,6 +526,7 @@ impl<'a> RecordBatchDecoder<'a> { projection: None, require_alignment: false, skip_validation: UnsafeFlag::new(), + custom_metadata: None, }) } @@ -544,6 +565,15 @@ 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: Option>, + ) -> 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 +584,10 @@ impl<'a> RecordBatchDecoder<'a> { .collect(); let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize)); + let custom_metadata = self.custom_metadata.take(); 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 +639,12 @@ impl<'a> RecordBatchDecoder<'a> { assert!(variadic_counts.is_empty()); RecordBatch::try_new_with_options(schema, children, &options) } - } + }; + + batch.map(|b| match custom_metadata { + Some(m) => b.with_custom_metadata(Arc::new(m)), + None => b, + }) } fn next_buffer(&mut self) -> Result { @@ -755,6 +791,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, @@ -1117,6 +1158,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 _), @@ -1128,6 +1170,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) } @@ -1682,6 +1725,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( @@ -1694,6 +1738,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 46e2dd7739e0..fdb8138902d1 100644 --- a/arrow-ipc/src/writer.rs +++ b/arrow-ipc/src/writer.rs @@ -591,12 +591,21 @@ 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() + .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); 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 +4329,216 @@ mod tests { let read_batch = reader.next().unwrap().unwrap(); assert_eq!(read_batch, batch2); } + + fn map_of>(items: I) -> CustomMetadata { + Arc::new(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)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap() + .with_custom_metadata(map_of([("key".to_string(), "value".to_string())])); + + // 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().and_then(|m| m.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(map_of([("key".to_string(), "value".to_string())])); + + // 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().and_then(|m| m.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(map_of([("key".to_string(), "value".to_string())])); + + // 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().and_then(|m| m.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(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(map_of([("batch_index".to_string(), "1".to_string())])); + + // 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() + .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")), + 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_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_none()); + 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() + .and_then(|m| m.get("batch_index")), + Some(&"0".to_string()) + ); + assert_eq!( + 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() + .and_then(|m| m.get("batch_index")), + Some(&"1".to_string()) + ); + assert_eq!( + batches[1].custom_metadata().and_then(|m| m.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().and_then(|m| m.get("batch_index")), + Some(&"0".to_string()) + ); + assert_eq!( + 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().and_then(|m| m.get("batch_index")), + Some(&"1".to_string()) + ); + assert_eq!( + batch1.custom_metadata().and_then(|m| m.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 000000000000..e7982d807ad6 Binary files /dev/null and b/arrow-ipc/test_data/custom_metadata.arrow_file differ 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 000000000000..b6b40a6c39e0 Binary files /dev/null and b/arrow-ipc/test_data/custom_metadata.arrow_stream differ diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index e95d01f2b592..ad7c5fdfed81 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().cloned(); unsafe { - Ok(RecordBatch::new_unchecked( - record_batch.schema(), - filtered_arrays, - self.count, - )) + let mut batch = + RecordBatch::new_unchecked(record_batch.schema(), filtered_arrays, self.count); + if let Some(m) = custom_metadata { + 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 cbb65ac915dd..67fe850f738c 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1116,7 +1116,11 @@ pub fn take_record_batch( .iter() .map(|c| take(c, indices, None)) .collect::, _>>()?; - RecordBatch::try_new(record_batch.schema(), columns) + 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(Arc::new(m)), + None => b, + }) } #[cfg(test)]