This is an automated email from the ASF dual-hosted git repository. CurtHagenlocher pushed a commit to branch ipc-message-custom-metadata in repository https://gitbox.apache.org/repos/asf/arrow-dotnet.git
commit 42a2b83f363b71752c8c3a9d54b668b7c814d023 Author: Christoph Mettler <[email protected]> AuthorDate: Sun Mar 8 20:14:29 2026 +0100 Add WriteRecordBatch overload with custom_metadata Adds WriteRecordBatch(batch, customMetadata) and its async counterpart to ArrowStreamWriter, allowing callers to attach per-message custom_metadata key-value pairs when writing IPC streams. The Arrow IPC flatbuf Message already defines a custom_metadata field, and pyarrow supports writing it via write_batch(batch, custom_metadata). This brings the C# writer to parity. Includes round-trip tests verifying custom_metadata survives write → read through ArrowStreamWriter/ArrowStreamReader. Co-Authored-By: Claude Opus 4.6 <[email protected]> --- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 57 +++++++- test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs | 159 ++++++++++++++++++++++ 2 files changed, 212 insertions(+), 4 deletions(-) diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index a39caa6..a13725f 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -806,6 +806,11 @@ namespace Apache.Arrow.Ipc } private protected void WriteRecordBatchInternal(RecordBatch recordBatch) + { + WriteRecordBatchInternal(recordBatch, customMetadata: null); + } + + private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata) { // TODO: Truncate buffers with extraneous padding / unused capacity @@ -829,6 +834,14 @@ namespace Apache.Arrow.Ipc VectorOffset buffersVectorOffset = Builder.EndVector(); + // Build custom metadata for the Message if provided + VectorOffset customMetadataVectorOffset = default; + if (customMetadata != null && customMetadata.Count > 0) + { + Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata); + customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + // Serialize record batch StartingWritingRecordBatch(); @@ -840,14 +853,21 @@ namespace Apache.Arrow.Ipc variadicCountsOffset); long metadataLength = WriteMessage(Flatbuf.MessageHeader.RecordBatch, - recordBatchOffset, recordBatchBuilder.TotalLength); + recordBatchOffset, recordBatchBuilder.TotalLength, customMetadataVectorOffset); long bufferLength = WriteBufferData(recordBatchBuilder.Buffers); FinishedWritingRecordBatch(bufferLength, metadataLength); } + private protected Task WriteRecordBatchInternalAsync(RecordBatch recordBatch, + CancellationToken cancellationToken = default) + { + return WriteRecordBatchInternalAsync(recordBatch, customMetadata: null, cancellationToken); + } + private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBatch, + IReadOnlyDictionary<string, string> customMetadata, CancellationToken cancellationToken = default) { if (!HasWrittenSchema) @@ -870,6 +890,14 @@ namespace Apache.Arrow.Ipc VectorOffset buffersVectorOffset = Builder.EndVector(); + // Build custom metadata for the Message if provided + VectorOffset customMetadataVectorOffset = default; + if (customMetadata != null && customMetadata.Count > 0) + { + Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata); + customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + // Serialize record batch StartingWritingRecordBatch(); @@ -882,6 +910,7 @@ namespace Apache.Arrow.Ipc long metadataLength = await WriteMessageAsync(Flatbuf.MessageHeader.RecordBatch, recordBatchOffset, recordBatchBuilder.TotalLength, + customMetadataVectorOffset, cancellationToken).ConfigureAwait(false); long bufferLength = await WriteBufferDataAsync(recordBatchBuilder.Buffers, cancellationToken).ConfigureAwait(false); @@ -1132,11 +1161,21 @@ namespace Apache.Arrow.Ipc WriteRecordBatchInternal(recordBatch); } + public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata) + { + WriteRecordBatchInternal(recordBatch, customMetadata); + } + public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default) { return WriteRecordBatchInternalAsync(recordBatch, cancellationToken); } + public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata, CancellationToken cancellationToken = default) + { + return WriteRecordBatchInternalAsync(recordBatch, customMetadata, cancellationToken); + } + public void WriteStart() { if (!HasWrittenStart) @@ -1347,12 +1386,13 @@ namespace Apache.Arrow.Ipc /// The number of bytes written to the stream. /// </returns> private protected long WriteMessage<T>( - Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength) + Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength, + VectorOffset customMetadataOffset = default) where T : struct { Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage( Builder, CurrentMetadataVersion, headerType, headerOffset.Value, - bodyLength); + bodyLength, customMetadataOffset); Builder.Finish(messageOffset.Value); @@ -1376,14 +1416,23 @@ namespace Apache.Arrow.Ipc /// <returns> /// The number of bytes written to the stream. /// </returns> + private protected virtual ValueTask<long> WriteMessageAsync<T>( + Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength, + CancellationToken cancellationToken) + where T : struct + { + return WriteMessageAsync(headerType, headerOffset, bodyLength, default, cancellationToken); + } + private protected virtual async ValueTask<long> WriteMessageAsync<T>( Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength, + VectorOffset customMetadataOffset, CancellationToken cancellationToken) where T : struct { Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage( Builder, CurrentMetadataVersion, headerType, headerOffset.Value, - bodyLength); + bodyLength, customMetadataOffset); Builder.Finish(messageOffset.Value); diff --git a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs index 1a4b5a6..1df27fd 100644 --- a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs @@ -736,5 +736,164 @@ namespace Apache.Arrow.Tests Assert.True(allocator.Statistics.Allocations > 0); Assert.Equal(0, allocator.Rented); } + + [Fact] + public void WriteCustomMetadata_RoundTrips() + { + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 10); + var customMetadata = new Dictionary<string, string> + { + ["rpc.method"] = "add", + ["rpc.version"] = "1", + ["request_id"] = "abc-123", + }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(originalBatch, customMetadata); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + RecordBatch readBatch = reader.ReadNextRecordBatch(); + Assert.NotNull(readBatch); + ArrowReaderVerifier.CompareBatches(originalBatch, readBatch); + + var readMetadata = reader.LastBatchCustomMetadata; + Assert.NotNull(readMetadata); + Assert.Equal(3, readMetadata.Count); + Assert.Equal("add", readMetadata["rpc.method"]); + Assert.Equal("1", readMetadata["rpc.version"]); + Assert.Equal("abc-123", readMetadata["request_id"]); + } + + [Fact] + public async Task WriteCustomMetadataAsync_RoundTrips() + { + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 10); + var customMetadata = new Dictionary<string, string> + { + ["key1"] = "value1", + ["key2"] = "value2", + }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + await writer.WriteRecordBatchAsync(originalBatch, customMetadata); + await writer.WriteEndAsync(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + RecordBatch readBatch = reader.ReadNextRecordBatch(); + Assert.NotNull(readBatch); + ArrowReaderVerifier.CompareBatches(originalBatch, readBatch); + + Assert.NotNull(reader.LastBatchCustomMetadata); + Assert.Equal("value1", reader.LastBatchCustomMetadata["key1"]); + Assert.Equal("value2", reader.LastBatchCustomMetadata["key2"]); + } + + [Fact] + public void WriteCustomMetadata_MultipleBatches_EachHasOwnMetadata() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta1 = new Dictionary<string, string> { ["batch"] = "first" }; + var meta2 = new Dictionary<string, string> { ["batch"] = "second", ["extra"] = "data" }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch, meta1); + writer.WriteRecordBatch(batch, meta2); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + + reader.ReadNextRecordBatch(); + Assert.NotNull(reader.LastBatchCustomMetadata); + Assert.Single(reader.LastBatchCustomMetadata); + Assert.Equal("first", reader.LastBatchCustomMetadata["batch"]); + + reader.ReadNextRecordBatch(); + Assert.NotNull(reader.LastBatchCustomMetadata); + Assert.Equal(2, reader.LastBatchCustomMetadata.Count); + Assert.Equal("second", reader.LastBatchCustomMetadata["batch"]); + Assert.Equal("data", reader.LastBatchCustomMetadata["extra"]); + } + + [Fact] + public void WriteWithoutCustomMetadata_LastBatchCustomMetadataIsNull() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + reader.ReadNextRecordBatch(); + Assert.Null(reader.LastBatchCustomMetadata); + } + + [Fact] + public void WriteCustomMetadata_MixedBatches_WithAndWithoutMetadata() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary<string, string> { ["key"] = "value" }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch, meta); + writer.WriteRecordBatch(batch); // no metadata + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + + reader.ReadNextRecordBatch(); + Assert.NotNull(reader.LastBatchCustomMetadata); + Assert.Equal("value", reader.LastBatchCustomMetadata["key"]); + + reader.ReadNextRecordBatch(); + Assert.Null(reader.LastBatchCustomMetadata); + } + + [Fact] + public void WriteCustomMetadata_EmptyValues_RoundTrips() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary<string, string> { ["empty"] = "" }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch, meta); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + reader.ReadNextRecordBatch(); + Assert.NotNull(reader.LastBatchCustomMetadata); + Assert.Equal("", reader.LastBatchCustomMetadata["empty"]); + } } }
