From 7a5bacf9b6df0a5d1588e38a1bd9e1418d4328b1 Mon Sep 17 00:00:00 2001 From: Christoph Mettler Date: Sun, 8 Mar 2026 20:14:21 +0100 Subject: [PATCH 1/6] Expose IPC Message custom_metadata on ArrowStreamReader The Arrow IPC format supports custom_metadata on each Message (RecordBatch), but the C# implementation currently ignores it on read. This adds a LastBatchCustomMetadata property to ArrowStreamReader that exposes the key-value pairs from the most recently read batch's Message. This is the read-side counterpart to pyarrow's read_next_batch_with_custom_metadata() and enables use cases like RPC frameworks that embed method routing or log metadata in per-batch custom_metadata fields. Co-Authored-By: Claude Opus 4.6 --- .../Ipc/ArrowReaderImplementation.cs | 23 +++++++++++++++++++ src/Apache.Arrow/Ipc/ArrowStreamReader.cs | 8 +++++++ 2 files changed, 31 insertions(+) diff --git a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs index 2c380e5d..acac52fd 100644 --- a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs +++ b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs @@ -81,6 +81,11 @@ protected virtual void Dispose(bool disposing) public abstract ValueTask ReadNextRecordBatchAsync(CancellationToken cancellationToken); public abstract RecordBatch ReadNextRecordBatch(); + /// + /// Custom metadata from the most recently read RecordBatch Message, if any. + /// + internal IReadOnlyDictionary LastBatchCustomMetadata { get; private protected set; } + internal static T ReadMessage(ByteBuffer bb) where T : struct, IFlatbufferObject { @@ -148,6 +153,7 @@ protected RecordBatch CreateArrowObjectFromMessage( } List arrays = BuildArrays(message.Version, Schema, bodyByteBuffer, rb); + LastBatchCustomMetadata = ReadMessageCustomMetadata(message); return new RecordBatch(Schema, memoryOwner, arrays, (int)rb.Length); default: // NOTE: Skip unsupported message type @@ -158,6 +164,23 @@ protected RecordBatch CreateArrowObjectFromMessage( return null; } + private static IReadOnlyDictionary ReadMessageCustomMetadata(Flatbuf.Message message) + { + int count = message.CustomMetadataLength; + if (count == 0) + return null; + + var result = new Dictionary(count); + for (int i = 0; i < count; i++) + { + Flatbuf.KeyValue kv = message.CustomMetadata(i).GetValueOrDefault(); + string key = kv.Key; + if (key != null) + result[key] = kv.Value ?? ""; + } + return result; + } + internal static ByteBuffer CreateByteBuffer(ReadOnlyMemory buffer) { return new ByteBuffer(new ReadOnlyMemoryBufferAllocator(buffer), 0); diff --git a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs index afa3713d..8379b12b 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs @@ -14,6 +14,7 @@ // limitations under the License. using System; +using System.Collections.Generic; using System.IO; using System.Threading; using System.Threading.Tasks; @@ -151,5 +152,12 @@ public RecordBatch ReadNextRecordBatch() { return _implementation.ReadNextRecordBatch(); } + + /// + /// Custom metadata from the most recently read RecordBatch Message. + /// Updated after each call to ReadNextRecordBatch/ReadNextRecordBatchAsync. + /// Returns null if the last batch had no custom metadata. + /// + public IReadOnlyDictionary LastBatchCustomMetadata => _implementation.LastBatchCustomMetadata; } } From 42a2b83f363b71752c8c3a9d54b668b7c814d023 Mon Sep 17 00:00:00 2001 From: Christoph Mettler Date: Sun, 8 Mar 2026 20:14:29 +0100 Subject: [PATCH 2/6] Add WriteRecordBatch overload with custom_metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 57 ++++++- .../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 a39caa66..a13725f0 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -806,6 +806,11 @@ public ArrowStreamWriter(Stream baseStream, Schema schema, bool leaveOpen, IpcOp } private protected void WriteRecordBatchInternal(RecordBatch recordBatch) + { + WriteRecordBatchInternal(recordBatch, customMetadata: null); + } + + private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOnlyDictionary customMetadata) { // TODO: Truncate buffers with extraneous padding / unused capacity @@ -829,6 +834,14 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch) VectorOffset buffersVectorOffset = Builder.EndVector(); + // Build custom metadata for the Message if provided + VectorOffset customMetadataVectorOffset = default; + if (customMetadata != null && customMetadata.Count > 0) + { + Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); + customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + // Serialize record batch StartingWritingRecordBatch(); @@ -840,14 +853,21 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch) 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 customMetadata, CancellationToken cancellationToken = default) { if (!HasWrittenSchema) @@ -870,6 +890,14 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat VectorOffset buffersVectorOffset = Builder.EndVector(); + // Build custom metadata for the Message if provided + VectorOffset customMetadataVectorOffset = default; + if (customMetadata != null && customMetadata.Count > 0) + { + Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); + customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + // Serialize record batch StartingWritingRecordBatch(); @@ -882,6 +910,7 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat 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 @@ public virtual void WriteRecordBatch(RecordBatch recordBatch) WriteRecordBatchInternal(recordBatch); } + public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionary customMetadata) + { + WriteRecordBatchInternal(recordBatch, customMetadata); + } + public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default) { return WriteRecordBatchInternalAsync(recordBatch, cancellationToken); } + public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, IReadOnlyDictionary customMetadata, CancellationToken cancellationToken = default) + { + return WriteRecordBatchInternalAsync(recordBatch, customMetadata, cancellationToken); + } + public void WriteStart() { if (!HasWrittenStart) @@ -1347,12 +1386,13 @@ await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, cancellat /// The number of bytes written to the stream. /// private protected long WriteMessage( - Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength) + Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength, + VectorOffset customMetadataOffset = default) where T : struct { Offset messageOffset = Flatbuf.Message.CreateMessage( Builder, CurrentMetadataVersion, headerType, headerOffset.Value, - bodyLength); + bodyLength, customMetadataOffset); Builder.Finish(messageOffset.Value); @@ -1376,14 +1416,23 @@ private protected long WriteMessage( /// /// The number of bytes written to the stream. /// + private protected virtual ValueTask WriteMessageAsync( + Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength, + CancellationToken cancellationToken) + where T : struct + { + return WriteMessageAsync(headerType, headerOffset, bodyLength, default, cancellationToken); + } + private protected virtual async ValueTask WriteMessageAsync( Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength, + VectorOffset customMetadataOffset, CancellationToken cancellationToken) where T : struct { Offset 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 1a4b5a6b..1df27fd5 100644 --- a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs @@ -736,5 +736,164 @@ public async Task MemoryOwnerDisposalSlicedArray(int sliceOffset, int sliceLengt 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 + { + ["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 + { + ["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 { ["batch"] = "first" }; + var meta2 = new Dictionary { ["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 { ["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 { ["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"]); + } } } From b68ca155cbf3a7a8e8ef06c7d01d5b04e5bdbce0 Mon Sep 17 00:00:00 2001 From: Christoph Mettler Date: Wed, 11 Mar 2026 20:17:09 +0100 Subject: [PATCH 3/6] =?UTF-8?q?Add=20custom=5Fmetadata=20cross-language=20?= =?UTF-8?q?tests=20using=20pythonnet=20+=20pyarrow:=20-=20C#=20writes=20IP?= =?UTF-8?q?C=20stream=20with=20custom=5Fmetadata=20=E2=86=92=20Python=20re?= =?UTF-8?q?ads=20via=20=20=20read=5Fnext=5Fbatch=5Fwith=5Fcustom=5Fmetadat?= =?UTF-8?q?a()=20-=20Python=20writes=20IPC=20stream=20with=20custom=5Fmeta?= =?UTF-8?q?data=20=E2=86=92=20C#=20reads=20via=20=20=20LastBatchCustomMeta?= =?UTF-8?q?data?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../CustomMetadataPythonTests.cs | 183 ++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs diff --git a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs new file mode 100644 index 00000000..56f68f32 --- /dev/null +++ b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs @@ -0,0 +1,183 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +using System; +using System.Collections.Generic; +using System.IO; +using Apache.Arrow.Ipc; +using Python.Runtime; +using Xunit; + +namespace Apache.Arrow.Tests +{ + + // ------------------------------------------------------------------- + // Cross-language Python tests for custom_metadata + // ------------------------------------------------------------------- + + public class CustomMetadataPythonTests : IClassFixture + { + public class PythonNet : IDisposable + { + public bool Initialized { get; } + + public bool VersionMismatch { get; } + + public PythonNet() + { + bool pythonSet = Environment.GetEnvironmentVariable("PYTHONNET_PYDLL") != null; + if (!pythonSet) + { + Initialized = false; + return; + } + + try + { + PythonEngine.Initialize(); + } + catch (NotSupportedException e) when (e.Message.Contains("Python ABI ") && e.Message.Contains("not supported")) + { + Initialized = false; + VersionMismatch = true; + return; + } + + if (System.Runtime.InteropServices.RuntimeInformation.IsOSPlatform(System.Runtime.InteropServices.OSPlatform.Windows) && + PythonEngine.PythonPath.IndexOf("dlls", StringComparison.OrdinalIgnoreCase) < 0) + { + dynamic sys = Py.Import("sys"); + sys.path.append(Path.Combine(Path.GetDirectoryName(Environment.GetEnvironmentVariable("PYTHONNET_PYDLL")), "DLLs")); + } + + Initialized = true; + } + + public void Dispose() + { + PythonEngine.Shutdown(); + } + } + + public CustomMetadataPythonTests(PythonNet pythonNet) + { + if (!pythonNet.Initialized) + { + var errorReason = pythonNet.VersionMismatch ? "Python version is incompatible with PythonNet" : "PYTHONNET_PYDLL not set"; + + bool inCIJob = Environment.GetEnvironmentVariable("GITHUB_ACTIONS") == "true"; + bool inVerificationJob = Environment.GetEnvironmentVariable("TEST_CSHARP") == "1"; + + Skip.If(inVerificationJob || !inCIJob, $"{errorReason}; skipping custom metadata Python tests."); + + throw new Exception($"{errorReason}; cannot run custom metadata Python tests."); + } + } + + // ------------------------------------------------------------------- + // C# writes IPC with custom_metadata → Python reads + // ------------------------------------------------------------------- + + [SkippableFact] + public void ExportCustomMetadata_PythonReads() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var batchMetadata = new Dictionary + { + ["rpc.method"] = "greet", + ["request_id"] = "abc-123", + ["custom_key"] = "custom_value", + }; + + // Serialize to IPC stream with custom batch metadata + byte[] ipcBytes; + using (var ms = new MemoryStream()) + { + using (var writer = new ArrowStreamWriter(ms, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch, batchMetadata); + writer.WriteEnd(); + } + ipcBytes = ms.ToArray(); + } + + // Python reads and verifies custom_metadata + using (Py.GIL()) + { + dynamic pa = Py.Import("pyarrow"); + dynamic reader = pa.ipc.open_stream(pa.BufferReader(ipcBytes.ToPython())); + + PyObject result = reader.read_next_batch_with_custom_metadata(); + dynamic pyBatch = result[0]; + dynamic customMeta = result[1]; + + // Verify batch data round-tripped + Assert.Equal(5, (int)pyBatch.num_rows); + + // Verify custom_metadata (pyarrow returns bytes — decode to str) + Assert.Equal("greet", (string)customMeta["rpc.method"].decode()); + Assert.Equal("abc-123", (string)customMeta["request_id"].decode()); + Assert.Equal("custom_value", (string)customMeta["custom_key"].decode()); + } + } + + // ------------------------------------------------------------------- + // Python writes IPC with custom_metadata → C# reads + // ------------------------------------------------------------------- + + [SkippableFact] + public void ImportCustomMetadata_PythonWrites() + { + byte[] ipcBytes; + + // Python creates a batch with custom_metadata and serializes to IPC + using (Py.GIL()) + { + dynamic pa = Py.Import("pyarrow"); + dynamic io = Py.Import("io"); + + dynamic pyBatch = pa.record_batch(new PyList(new PyObject[] + { + pa.array(new int[] { 1, 2, 3, 4, 5 }), + }), new[] { "x" }); + + dynamic buf = io.BytesIO(); + dynamic writer = pa.ipc.new_stream(buf, pyBatch.schema); + dynamic customMeta = pa.KeyValueMetadata(new PyDict + { + ["origin"] = "python".ToPython(), + ["version"] = "2".ToPython(), + }); + writer.write_batch(pyBatch, custom_metadata: customMeta); + writer.close(); + + ipcBytes = ((PyObject)buf.getvalue()).As(); + } + + // C# reads and verifies custom_metadata + using var ms = new MemoryStream(ipcBytes); + using var reader = new ArrowStreamReader(ms); + + RecordBatch batch = reader.ReadNextRecordBatch(); + Assert.NotNull(batch); + Assert.Equal(5, batch.Length); + + var metadata = reader.LastBatchCustomMetadata; + Assert.NotNull(metadata); + Assert.Equal("python", metadata["origin"]); + Assert.Equal("2", metadata["version"]); + } + } +} From fcff0b807133927e9cd9c27e31477aab34645c6b Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 24 Aug 2026 15:25:36 -0400 Subject: [PATCH 4/6] Address Copilot review comments on #424 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - FlightDataStream: override the customMetadata-aware WriteMessageAsync overload (the one WriteRecordBatchInternalAsync now actually calls) instead of the old 4-arg overload, so Flight writes aren't silently routed through the base implementation and don't bypass DataHeader capture. - ArrowStreamWriter: validate that caller-supplied custom metadata has no null keys/values before building FlatBuffer offsets, so failures are reported as a clear ArgumentException rather than an opaque FlatBufferBuilder exception. - ArrowStreamReader: correct the LastBatchCustomMetadata XML doc to describe its actual update semantics (it's left unchanged when a read call returns null, e.g. at end of stream). - CustomMetadataPythonTests: use the repo's shared PythonNetFixture + [Collection("PythonNet")] instead of a private per-class Python.NET init/shutdown, avoiding double-Initialize/premature-Shutdown races with other Python.NET tests; this also fixes the missing Py.GIL() guard around the Windows sys.path append, since the shared fixture already wraps that in using (Py.GIL()). Verified with: - dotnet build Apache.Arrow.sln — 0 warnings/errors - dotnet test test/Apache.Arrow.Tests/Apache.Arrow.Tests.csproj — 1876 passed, 30 skipped, 0 failed - dotnet format Apache.Arrow.sln --exclude src/Apache.Arrow/Flatbuf/FlatBuffers/ --verify-no-changes — clean --- .../Internal/FlightDataStream.cs | 4 +- src/Apache.Arrow/Ipc/ArrowStreamReader.cs | 6 +- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 21 +++++++ .../CustomMetadataPythonTests.cs | 59 ++----------------- 4 files changed, 31 insertions(+), 59 deletions(-) diff --git a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs index 50b2a40b..3ab43318 100644 --- a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs +++ b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs @@ -91,11 +91,11 @@ public async Task Write(RecordBatch recordBatch, ByteString applicationMetadata) await _clientStreamWriter.WriteAsync(_currentFlightData).ConfigureAwait(false); } - private protected override ValueTask WriteMessageAsync(MessageHeader headerType, Offset headerOffset, int bodyLength, CancellationToken cancellationToken) + private protected override ValueTask WriteMessageAsync(MessageHeader headerType, Offset headerOffset, int bodyLength, VectorOffset customMetadataOffset, CancellationToken cancellationToken) { Offset messageOffset = Flatbuf.Message.CreateMessage( Builder, CurrentMetadataVersion, headerType, headerOffset.Value, - bodyLength); + bodyLength, customMetadataOffset); Builder.Finish(messageOffset.Value); diff --git a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs index 8379b12b..e7b5861e 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs @@ -155,8 +155,10 @@ public RecordBatch ReadNextRecordBatch() /// /// Custom metadata from the most recently read RecordBatch Message. - /// Updated after each call to ReadNextRecordBatch/ReadNextRecordBatchAsync. - /// Returns null if the last batch had no custom metadata. + /// Set whenever ReadNextRecordBatch/ReadNextRecordBatchAsync successfully reads a + /// RecordBatch message; left unchanged when a call returns null (e.g. at the end of + /// the stream), so it continues to reflect the last RecordBatch that was read. + /// Returns null if that batch had no custom metadata. /// public IReadOnlyDictionary LastBatchCustomMetadata => _implementation.LastBatchCustomMetadata; } diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index a13725f0..c8e60323 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -838,6 +838,7 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOn VectorOffset customMetadataVectorOffset = default; if (customMetadata != null && customMetadata.Count > 0) { + ValidateCustomMetadata(customMetadata); Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); } @@ -894,6 +895,7 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat VectorOffset customMetadataVectorOffset = default; if (customMetadata != null && customMetadata.Count > 0) { + ValidateCustomMetadata(customMetadata); Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); } @@ -1330,6 +1332,25 @@ private VectorOffset GetFieldMetadataOffset(Field field) return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered); } + /// + /// Validates that a caller-supplied custom metadata dictionary contains no null keys or values, + /// so that failures are reported clearly rather than as an opaque exception from the FlatBuffer builder. + /// + private static void ValidateCustomMetadata(IReadOnlyDictionary customMetadata) + { + foreach (KeyValuePair metadatum in customMetadata) + { + if (metadatum.Key == null) + { + throw new ArgumentException("Custom metadata must not contain null keys.", nameof(customMetadata)); + } + if (metadatum.Value == null) + { + throw new ArgumentException($"Custom metadata value for key '{metadatum.Key}' must not be null.", nameof(customMetadata)); + } + } + } + private Offset[] GetMetadataOffsets(IReadOnlyDictionary metadata) { Debug.Assert(metadata != null); diff --git a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs index 56f68f32..22317231 100644 --- a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs +++ b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs @@ -27,63 +27,12 @@ namespace Apache.Arrow.Tests // Cross-language Python tests for custom_metadata // ------------------------------------------------------------------- - public class CustomMetadataPythonTests : IClassFixture + [Collection("PythonNet")] + public class CustomMetadataPythonTests { - public class PythonNet : IDisposable + public CustomMetadataPythonTests(PythonNetFixture pythonNet) { - public bool Initialized { get; } - - public bool VersionMismatch { get; } - - public PythonNet() - { - bool pythonSet = Environment.GetEnvironmentVariable("PYTHONNET_PYDLL") != null; - if (!pythonSet) - { - Initialized = false; - return; - } - - try - { - PythonEngine.Initialize(); - } - catch (NotSupportedException e) when (e.Message.Contains("Python ABI ") && e.Message.Contains("not supported")) - { - Initialized = false; - VersionMismatch = true; - return; - } - - if (System.Runtime.InteropServices.RuntimeInformation.IsOSPlatform(System.Runtime.InteropServices.OSPlatform.Windows) && - PythonEngine.PythonPath.IndexOf("dlls", StringComparison.OrdinalIgnoreCase) < 0) - { - dynamic sys = Py.Import("sys"); - sys.path.append(Path.Combine(Path.GetDirectoryName(Environment.GetEnvironmentVariable("PYTHONNET_PYDLL")), "DLLs")); - } - - Initialized = true; - } - - public void Dispose() - { - PythonEngine.Shutdown(); - } - } - - public CustomMetadataPythonTests(PythonNet pythonNet) - { - if (!pythonNet.Initialized) - { - var errorReason = pythonNet.VersionMismatch ? "Python version is incompatible with PythonNet" : "PYTHONNET_PYDLL not set"; - - bool inCIJob = Environment.GetEnvironmentVariable("GITHUB_ACTIONS") == "true"; - bool inVerificationJob = Environment.GetEnvironmentVariable("TEST_CSHARP") == "1"; - - Skip.If(inVerificationJob || !inCIJob, $"{errorReason}; skipping custom metadata Python tests."); - - throw new Exception($"{errorReason}; cannot run custom metadata Python tests."); - } + pythonNet.EnsureInitialized(); } // ------------------------------------------------------------------- From 284d3003ceb3580a25b6a0ac5357f813b0fdcd32 Mon Sep 17 00:00:00 2001 From: Curt Hagenlocher Date: Sun, 6 Sep 2026 21:06:42 -0700 Subject: [PATCH 5/6] Address review feedback: make the "start" preamble unskippable The new WriteRecordBatch(batch, customMetadata) overloads went straight to WriteRecordBatchInternal, but ArrowFileWriter relied on overriding each public WriteRecordBatch to call WriteStart() first. Calling the new overload on an ArrowFileWriter therefore skipped the ARROW1 file magic and silently produced a file that ArrowFileReader rejects with "Invalid magic at offset <6>". Rather than adding two more overrides that a future overload could again forget, move the WriteStart()/WriteStartAsync() call into WriteRecordBatchInternal, where every write path must pass through it. Both are idempotent, so the byte output is unchanged for the stream writer, the file writer and Flight. ArrowFileWriter's WriteRecordBatch/WriteRecordBatchAsync overrides are now redundant and removed. Also: - Remove the second virtual WriteMessageAsync overload. Two virtual overloads where one forwards to the other is the trap that already routed Flight's record batch writes past FlightDataStream's override; the sync WriteMessage has always been a single method with a defaulted customMetadataOffset. Callers that do not supply metadata now pass default explicitly. - Remove the private protected WriteRecordBatchInternal/WriteRecordBatchInternalAsync forwarding overloads. They carry no backwards-compatibility obligation, and fewer near-identical overloads means fewer places for the metadata argument to get silently dropped. - Validate custom metadata before anything is written instead of part-way through building the message, so a rejected dictionary leaves the writer usable, and hoist the duplicated offset-building block into GetCustomMetadataOffset. - Read Message.custom_metadata the same way schema and field metadata are already read in MessageSerializer, rather than skipping null keys and rewriting null values as "". Tests: ArrowFileWriter round-trips with custom metadata (sync and async) and still emits the file magic, custom metadata after an explicit WriteStart, empty dictionary, null key and null value rejection, and writer reuse after a rejected dictionary. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01XhMo3XSWYHo1apHd9PzZTb --- .../Internal/FlightDataStream.cs | 4 +- src/Apache.Arrow/Ipc/ArrowFileWriter.cs | 21 ---- .../Ipc/ArrowReaderImplementation.cs | 19 ++- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 82 ++++++------- .../ArrowFileWriterTests.cs | 108 ++++++++++++++++++ .../ArrowStreamWriterTests.cs | 104 +++++++++++++++++ 6 files changed, 264 insertions(+), 74 deletions(-) diff --git a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs index 3ab43318..46a7ac20 100644 --- a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs +++ b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs @@ -55,7 +55,7 @@ public async Task SendSchema() var offset = SerializeSchema(Schema); CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - await WriteMessageAsync(MessageHeader.Schema, offset, 0, cancellationTokenSource.Token).ConfigureAwait(false); + await WriteMessageAsync(MessageHeader.Schema, offset, 0, default, cancellationTokenSource.Token).ConfigureAwait(false); await _clientStreamWriter.WriteAsync(_currentFlightData).ConfigureAwait(false); HasWrittenSchema = true; } @@ -81,7 +81,7 @@ public async Task Write(RecordBatch recordBatch, ByteString applicationMetadata) _currentFlightData.AppMetadata = applicationMetadata; } - await WriteRecordBatchInternalAsync(recordBatch).ConfigureAwait(false); + await WriteRecordBatchInternalAsync(recordBatch, customMetadata: null).ConfigureAwait(false); //Reset stream position this.BaseStream.Position = 0; diff --git a/src/Apache.Arrow/Ipc/ArrowFileWriter.cs b/src/Apache.Arrow/Ipc/ArrowFileWriter.cs index 91b7c298..cfdf2269 100644 --- a/src/Apache.Arrow/Ipc/ArrowFileWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowFileWriter.cs @@ -66,27 +66,6 @@ public ArrowFileWriter(Stream stream, Schema schema, bool leaveOpen, IpcOptions RecordBatchBlocks = new List(); } - public override void WriteRecordBatch(RecordBatch recordBatch) - { - // TODO: Compare record batch schema - - WriteStart(); - - WriteRecordBatchInternal(recordBatch); - } - - public override async Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default) - { - // TODO: Compare record batch schema - - await WriteStartAsync(cancellationToken).ConfigureAwait(false); - - cancellationToken.ThrowIfCancellationRequested(); - - await WriteRecordBatchInternalAsync(recordBatch, cancellationToken) - .ConfigureAwait(false); - } - private protected override void StartingWritingRecordBatch() { _currentRecordBatchOffset = BaseStream.Position; diff --git a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs index acac52fd..45fd7920 100644 --- a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs +++ b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs @@ -166,19 +166,16 @@ protected RecordBatch CreateArrowObjectFromMessage( private static IReadOnlyDictionary ReadMessageCustomMetadata(Flatbuf.Message message) { - int count = message.CustomMetadataLength; - if (count == 0) - return null; - - var result = new Dictionary(count); - for (int i = 0; i < count; i++) + Dictionary metadata = message.CustomMetadataLength > 0 + ? new Dictionary(message.CustomMetadataLength) : null; + for (int i = 0; i < message.CustomMetadataLength; i++) { - Flatbuf.KeyValue kv = message.CustomMetadata(i).GetValueOrDefault(); - string key = kv.Key; - if (key != null) - result[key] = kv.Value ?? ""; + Flatbuf.KeyValue keyValue = message.CustomMetadata(i).GetValueOrDefault(); + + metadata[keyValue.Key] = keyValue.Value; } - return result; + + return metadata; } internal static ByteBuffer CreateByteBuffer(ReadOnlyMemory buffer) diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index c8e60323..aa63985e 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -805,14 +805,17 @@ public ArrowStreamWriter(Stream baseStream, Schema schema, bool leaveOpen, IpcOp Builder, compressionType, Flatbuf.BodyCompressionMethod.BUFFER); } - private protected void WriteRecordBatchInternal(RecordBatch recordBatch) - { - WriteRecordBatchInternal(recordBatch, customMetadata: null); - } - private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOnlyDictionary customMetadata) { // TODO: Truncate buffers with extraneous padding / unused capacity + // TODO: Compare record batch schema + + ValidateCustomMetadata(customMetadata); + + // Derived writers use WriteStartInternal to emit a preamble before any message + // (ArrowFileWriter writes the file magic there). Doing this here rather than in + // the public entry points means a new WriteRecordBatch overload cannot skip it. + WriteStart(); if (!HasWrittenSchema) { @@ -834,14 +837,7 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOn VectorOffset buffersVectorOffset = Builder.EndVector(); - // Build custom metadata for the Message if provided - VectorOffset customMetadataVectorOffset = default; - if (customMetadata != null && customMetadata.Count > 0) - { - ValidateCustomMetadata(customMetadata); - Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); - customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); - } + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); // Serialize record batch @@ -861,16 +857,17 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch, IReadOn 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 customMetadata, CancellationToken cancellationToken = default) { + // TODO: Compare record batch schema + + ValidateCustomMetadata(customMetadata); + + // See the comment in WriteRecordBatchInternal. + await WriteStartAsync(cancellationToken).ConfigureAwait(false); + if (!HasWrittenSchema) { await WriteSchemaAsync(Schema, cancellationToken).ConfigureAwait(false); @@ -891,14 +888,7 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat VectorOffset buffersVectorOffset = Builder.EndVector(); - // Build custom metadata for the Message if provided - VectorOffset customMetadataVectorOffset = default; - if (customMetadata != null && customMetadata.Count > 0) - { - ValidateCustomMetadata(customMetadata); - Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); - customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); - } + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); // Serialize record batch @@ -1090,7 +1080,7 @@ private protected async Task WriteDictionaryAsync(long id, IArrowType valueType, using var builder = recordBatchBuilder; long metadataLength = await WriteMessageAsync(Flatbuf.MessageHeader.DictionaryBatch, - dictionaryBatchOffset, recordBatchBuilder.TotalLength, cancellationToken).ConfigureAwait(false); + dictionaryBatchOffset, recordBatchBuilder.TotalLength, default, cancellationToken).ConfigureAwait(false); long bufferLength = await WriteBufferDataAsync(recordBatchBuilder.Buffers, cancellationToken).ConfigureAwait(false); @@ -1160,7 +1150,7 @@ private protected virtual void FinishedWritingRecordBatch(long bodyLength, long public virtual void WriteRecordBatch(RecordBatch recordBatch) { - WriteRecordBatchInternal(recordBatch); + WriteRecordBatchInternal(recordBatch, customMetadata: null); } public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionary customMetadata) @@ -1170,7 +1160,7 @@ public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionar public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default) { - return WriteRecordBatchInternalAsync(recordBatch, cancellationToken); + return WriteRecordBatchInternalAsync(recordBatch, customMetadata: null, cancellationToken); } public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, IReadOnlyDictionary customMetadata, CancellationToken cancellationToken = default) @@ -1332,12 +1322,32 @@ private VectorOffset GetFieldMetadataOffset(Field field) return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered); } + /// + /// Builds the Message-level custom_metadata vector, or a default offset when there is none. + /// + private VectorOffset GetCustomMetadataOffset(IReadOnlyDictionary customMetadata) + { + if (customMetadata == null || customMetadata.Count == 0) + { + return default; + } + + Offset[] metadataOffsets = GetMetadataOffsets(customMetadata); + return Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + /// /// Validates that a caller-supplied custom metadata dictionary contains no null keys or values, - /// so that failures are reported clearly rather than as an opaque exception from the FlatBuffer builder. + /// so that failures are reported before anything is written rather than as an opaque exception + /// from the FlatBuffer builder part-way through a message. /// private static void ValidateCustomMetadata(IReadOnlyDictionary customMetadata) { + if (customMetadata == null) + { + return; + } + foreach (KeyValuePair metadatum in customMetadata) { if (metadatum.Key == null) @@ -1394,7 +1404,7 @@ private static void ValidateCustomMetadata(IReadOnlyDictionary c // Build message - await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, cancellationToken) + await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, default, cancellationToken) .ConfigureAwait(false); return schemaOffset; @@ -1437,14 +1447,6 @@ private protected long WriteMessage( /// /// The number of bytes written to the stream. /// - private protected virtual ValueTask WriteMessageAsync( - Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength, - CancellationToken cancellationToken) - where T : struct - { - return WriteMessageAsync(headerType, headerOffset, bodyLength, default, cancellationToken); - } - private protected virtual async ValueTask WriteMessageAsync( Flatbuf.MessageHeader headerType, Offset headerOffset, int bodyLength, VectorOffset customMetadataOffset, diff --git a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs index d810a53b..af5eee8c 100644 --- a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs @@ -311,6 +311,114 @@ public async Task WriteListArrayWithEmptyOffsets() await ValidateRecordBatchFile(stream, recordBatch, strictCompare: false); } + [Fact] + public void WriteCustomMetadata_StillWritesFileMagic() + { + // ArrowFileWriter has to emit the file magic before any message. Regression test for + // a WriteRecordBatch overload reaching WriteRecordBatchInternal without it. + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100); + var customMetadata = new Dictionary { ["batch"] = "first" }; + + var stream = new MemoryStream(); + using (var writer = new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(originalBatch, customMetadata); + writer.WriteEnd(); + } + + Assert.Equal( + ArrowFileConstants.Magic, + stream.ToArray().AsSpan(0, ArrowFileConstants.Magic.Length).ToArray()); + } + + [Fact] + public async Task WriteCustomMetadataAsync_StillWritesFileMagic() + { + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100); + var customMetadata = new Dictionary { ["batch"] = "first" }; + + var stream = new MemoryStream(); + using (var writer = new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + await writer.WriteRecordBatchAsync(originalBatch, customMetadata); + await writer.WriteEndAsync(); + } + + Assert.Equal( + ArrowFileConstants.Magic, + stream.ToArray().AsSpan(0, ArrowFileConstants.Magic.Length).ToArray()); + } + + [Fact] + public async Task WriteCustomMetadata_RoundTrips() + { + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100); + var customMetadata = new Dictionary + { + ["rpc.method"] = "add", + ["request_id"] = "abc-123", + }; + + var stream = new MemoryStream(); + using (var writer = new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(originalBatch, customMetadata); + writer.WriteEnd(); + } + + stream.Position = 0; + + await ValidateRecordBatchFile(stream, originalBatch); + + stream.Position = 0; + using var reader = new ArrowFileReader(stream); + Assert.NotNull(reader.ReadNextRecordBatch()); + Assert.Equal(customMetadata, reader.LastBatchCustomMetadata); + } + + [Fact] + public async Task WriteCustomMetadataAsync_RoundTrips() + { + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100); + var customMetadata = new Dictionary { ["key1"] = "value1" }; + + var stream = new MemoryStream(); + using (var writer = new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + await writer.WriteRecordBatchAsync(originalBatch, customMetadata); + await writer.WriteEndAsync(); + } + + stream.Position = 0; + + await ValidateRecordBatchFile(stream, originalBatch); + + stream.Position = 0; + using var reader = new ArrowFileReader(stream); + Assert.NotNull(await reader.ReadNextRecordBatchAsync()); + Assert.Equal(customMetadata, reader.LastBatchCustomMetadata); + } + + [Fact] + public async Task WriteCustomMetadata_AfterExplicitWriteStart_RoundTrips() + { + // WriteStart is idempotent, so writing it up front must not produce a second preamble. + RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100); + var customMetadata = new Dictionary { ["key1"] = "value1" }; + + var stream = new MemoryStream(); + using (var writer = new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true)) + { + writer.WriteStart(); + writer.WriteRecordBatch(originalBatch, customMetadata); + writer.WriteEnd(); + } + + stream.Position = 0; + + await ValidateRecordBatchFile(stream, originalBatch); + } + private static void Shuffle(int[] values, Random random) { var length = values.Length; diff --git a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs index 1df27fd5..5246e1a1 100644 --- a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs @@ -15,6 +15,7 @@ using System; using System.Buffers.Binary; +using System.Collections; using System.Collections.Generic; using System.IO; using System.Linq; @@ -875,6 +876,90 @@ public void WriteCustomMetadata_MixedBatches_WithAndWithoutMetadata() Assert.Null(reader.LastBatchCustomMetadata); } + [Fact] + public void WriteCustomMetadata_EmptyDictionary_WritesNoMetadata() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + writer.WriteRecordBatch(batch, new Dictionary()); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + reader.ReadNextRecordBatch(); + Assert.Null(reader.LastBatchCustomMetadata); + } + + [Fact] + public void WriteCustomMetadata_NullKey_Throws() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + // Dictionary rejects a null key, so go through a map that allows one. + var withNullKey = new NullTolerantMetadata(new KeyValuePair(null, "value")); + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + Assert.Throws(() => writer.WriteRecordBatch(batch, withNullKey)); + } + + [Fact] + public void WriteCustomMetadata_NullValue_Throws() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary { ["key"] = null }; + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + Assert.Throws(() => writer.WriteRecordBatch(batch, meta)); + } + + [Fact] + public async Task WriteCustomMetadataAsync_NullValue_Throws() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary { ["key"] = null }; + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + await Assert.ThrowsAsync( + () => writer.WriteRecordBatchAsync(batch, meta)); + } + + [Fact] + public void WriteCustomMetadata_RejectedMetadata_LeavesWriterUsable() + { + // Validation happens before anything is written, so a rejected dictionary must not + // leave the writer part-way through a message. + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var good = new Dictionary { ["key"] = "value" }; + var bad = new Dictionary { ["key"] = null }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + Assert.Throws(() => writer.WriteRecordBatch(batch, bad)); + writer.WriteRecordBatch(batch, good); + writer.WriteEnd(); + } + + stream.Position = 0; + + using var reader = new ArrowStreamReader(stream); + RecordBatch readBatch = reader.ReadNextRecordBatch(); + Assert.NotNull(readBatch); + ArrowReaderVerifier.CompareBatches(batch, readBatch); + Assert.Equal(good, reader.LastBatchCustomMetadata); + Assert.Null(reader.ReadNextRecordBatch()); + } + [Fact] public void WriteCustomMetadata_EmptyValues_RoundTrips() { @@ -895,5 +980,24 @@ public void WriteCustomMetadata_EmptyValues_RoundTrips() Assert.NotNull(reader.LastBatchCustomMetadata); Assert.Equal("", reader.LastBatchCustomMetadata["empty"]); } + + /// + /// A metadata collection that can hold a null key, which cannot. + /// + private sealed class NullTolerantMetadata : IReadOnlyDictionary + { + private readonly KeyValuePair[] _entries; + + public NullTolerantMetadata(params KeyValuePair[] entries) => _entries = entries; + + public int Count => _entries.Length; + public IEnumerable Keys => _entries.Select(e => e.Key); + public IEnumerable Values => _entries.Select(e => e.Value); + public string this[string key] => throw new NotSupportedException(); + public bool ContainsKey(string key) => throw new NotSupportedException(); + public bool TryGetValue(string key, out string value) => throw new NotSupportedException(); + public IEnumerator> GetEnumerator() => ((IEnumerable>)_entries).GetEnumerator(); + IEnumerator IEnumerable.GetEnumerator() => _entries.GetEnumerator(); + } } } From 304d018488c029586191447b17031377959514a6 Mon Sep 17 00:00:00 2001 From: Curt Hagenlocher Date: Mon, 7 Sep 2026 13:25:37 -0700 Subject: [PATCH 6/6] Return custom metadata with the batch instead of exposing reader state Replace the ArrowStreamReader.LastBatchCustomMetadata property with a RecordBatchWithMetadata result type, mirroring pyarrow's read_next_batch_with_custom_metadata() and the equivalent Arrow C++ struct: RecordBatchWithMetadata ReadNextRecordBatchWithCustomMetadata(); ValueTask ReadNextRecordBatchWithCustomMetadataAsync(...); A property that has to be read at exactly the right moment is easy to get out of step with the batch in hand, and it had no sensible value at the end of the stream. Pairing the two in the return value removes both problems and reads the same in the sync and async APIs. The struct deconstructs, so callers who want the pair can write `var (batch, metadata) = ...`. ArrowFileReader gains ReadRecordBatchWithCustomMetadataAsync(int index) so the indexed read has the same capability as the sequential one. The transient state on ArrowReaderImplementation stays, but it is internal and consumed immediately by the two new methods rather than being public surface. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01XhMo3XSWYHo1apHd9PzZTb --- src/Apache.Arrow/Ipc/ArrowFileReader.cs | 11 +++ src/Apache.Arrow/Ipc/ArrowStreamReader.cs | 37 +++++++-- .../Ipc/RecordBatchWithMetadata.cs | 49 ++++++++++++ .../ArrowFileWriterTests.cs | 15 +++- .../ArrowStreamWriterTests.cs | 77 +++++++++---------- .../CustomMetadataPythonTests.cs | 4 +- 6 files changed, 141 insertions(+), 52 deletions(-) create mode 100644 src/Apache.Arrow/Ipc/RecordBatchWithMetadata.cs diff --git a/src/Apache.Arrow/Ipc/ArrowFileReader.cs b/src/Apache.Arrow/Ipc/ArrowFileReader.cs index fa3f84a5..c8b63efb 100644 --- a/src/Apache.Arrow/Ipc/ArrowFileReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowFileReader.cs @@ -85,5 +85,16 @@ public ValueTask ReadRecordBatchAsync(int index, CancellationToken { return Implementation.ReadRecordBatchAsync(index, cancellationToken); } + + /// + /// Reads the record batch at the given index together with the custom metadata on its + /// IPC Message, which is null if the message carried none. + /// + public async ValueTask ReadRecordBatchWithCustomMetadataAsync(int index, CancellationToken cancellationToken = default) + { + RecordBatch batch = await Implementation.ReadRecordBatchAsync(index, cancellationToken).ConfigureAwait(false); + + return batch == null ? default : new RecordBatchWithMetadata(batch, Implementation.LastBatchCustomMetadata); + } } } diff --git a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs index e7b5861e..bdc1fb7f 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs @@ -154,12 +154,37 @@ public RecordBatch ReadNextRecordBatch() } /// - /// Custom metadata from the most recently read RecordBatch Message. - /// Set whenever ReadNextRecordBatch/ReadNextRecordBatchAsync successfully reads a - /// RecordBatch message; left unchanged when a call returns null (e.g. at the end of - /// the stream), so it continues to reflect the last RecordBatch that was read. - /// Returns null if that batch had no custom metadata. + /// Reads the next record batch together with the custom metadata on its IPC Message, + /// the counterpart of . /// - public IReadOnlyDictionary LastBatchCustomMetadata => _implementation.LastBatchCustomMetadata; + /// + /// The record batch and its custom metadata. At the end of the stream both + /// and + /// are null; the metadata is also + /// null for a batch whose message carried none. + /// + public async ValueTask ReadNextRecordBatchWithCustomMetadataAsync(CancellationToken cancellationToken = default) + { + RecordBatch batch = await _implementation.ReadNextRecordBatchAsync(cancellationToken).ConfigureAwait(false); + + return batch == null ? default : new RecordBatchWithMetadata(batch, _implementation.LastBatchCustomMetadata); + } + + /// + /// Reads the next record batch together with the custom metadata on its IPC Message, + /// the counterpart of . + /// + /// + /// The record batch and its custom metadata. At the end of the stream both + /// and + /// are null; the metadata is also + /// null for a batch whose message carried none. + /// + public RecordBatchWithMetadata ReadNextRecordBatchWithCustomMetadata() + { + RecordBatch batch = _implementation.ReadNextRecordBatch(); + + return batch == null ? default : new RecordBatchWithMetadata(batch, _implementation.LastBatchCustomMetadata); + } } } diff --git a/src/Apache.Arrow/Ipc/RecordBatchWithMetadata.cs b/src/Apache.Arrow/Ipc/RecordBatchWithMetadata.cs new file mode 100644 index 00000000..54e21033 --- /dev/null +++ b/src/Apache.Arrow/Ipc/RecordBatchWithMetadata.cs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +using System.Collections.Generic; + +namespace Apache.Arrow.Ipc +{ + /// + /// A record batch read from an Arrow IPC source, together with the custom metadata + /// carried on the IPC Message that held it. + /// + public readonly struct RecordBatchWithMetadata + { + public RecordBatchWithMetadata(RecordBatch batch, IReadOnlyDictionary customMetadata) + { + Batch = batch; + CustomMetadata = customMetadata; + } + + /// + /// The record batch that was read, or null at the end of the stream. + /// + public RecordBatch Batch { get; } + + /// + /// The Message-level custom metadata accompanying , or null if the + /// message carried none. + /// + public IReadOnlyDictionary CustomMetadata { get; } + + public void Deconstruct(out RecordBatch batch, out IReadOnlyDictionary customMetadata) + { + batch = Batch; + customMetadata = CustomMetadata; + } + } +} diff --git a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs index af5eee8c..f3b0a343 100644 --- a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs @@ -372,8 +372,14 @@ public async Task WriteCustomMetadata_RoundTrips() stream.Position = 0; using var reader = new ArrowFileReader(stream); - Assert.NotNull(reader.ReadNextRecordBatch()); - Assert.Equal(customMetadata, reader.LastBatchCustomMetadata); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + Assert.Equal(customMetadata, read.CustomMetadata); + + // The indexed read on ArrowFileReader reports the same metadata. + RecordBatchWithMetadata indexed = await reader.ReadRecordBatchWithCustomMetadataAsync(0); + Assert.NotNull(indexed.Batch); + Assert.Equal(customMetadata, indexed.CustomMetadata); } [Fact] @@ -395,8 +401,9 @@ public async Task WriteCustomMetadataAsync_RoundTrips() stream.Position = 0; using var reader = new ArrowFileReader(stream); - Assert.NotNull(await reader.ReadNextRecordBatchAsync()); - Assert.Equal(customMetadata, reader.LastBatchCustomMetadata); + RecordBatchWithMetadata read = await reader.ReadNextRecordBatchWithCustomMetadataAsync(); + Assert.NotNull(read.Batch); + Assert.Equal(customMetadata, read.CustomMetadata); } [Fact] diff --git a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs index 5246e1a1..172f69b7 100644 --- a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs @@ -759,16 +759,15 @@ public void WriteCustomMetadata_RoundTrips() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - RecordBatch readBatch = reader.ReadNextRecordBatch(); - Assert.NotNull(readBatch); - ArrowReaderVerifier.CompareBatches(originalBatch, readBatch); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + ArrowReaderVerifier.CompareBatches(originalBatch, read.Batch); - 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"]); + Assert.NotNull(read.CustomMetadata); + Assert.Equal(3, read.CustomMetadata.Count); + Assert.Equal("add", read.CustomMetadata["rpc.method"]); + Assert.Equal("1", read.CustomMetadata["rpc.version"]); + Assert.Equal("abc-123", read.CustomMetadata["request_id"]); } [Fact] @@ -791,13 +790,12 @@ public async Task WriteCustomMetadataAsync_RoundTrips() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - RecordBatch readBatch = reader.ReadNextRecordBatch(); + (RecordBatch readBatch, IReadOnlyDictionary readMetadata) = + await reader.ReadNextRecordBatchWithCustomMetadataAsync(); Assert.NotNull(readBatch); ArrowReaderVerifier.CompareBatches(originalBatch, readBatch); - Assert.NotNull(reader.LastBatchCustomMetadata); - Assert.Equal("value1", reader.LastBatchCustomMetadata["key1"]); - Assert.Equal("value2", reader.LastBatchCustomMetadata["key2"]); + Assert.Equal(customMetadata, readMetadata); } [Fact] @@ -819,20 +817,12 @@ public void WriteCustomMetadata_MultipleBatches_EachHasOwnMetadata() 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"]); + Assert.Equal(meta1, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); + Assert.Equal(meta2, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); } [Fact] - public void WriteWithoutCustomMetadata_LastBatchCustomMetadataIsNull() + public void WriteWithoutCustomMetadata_CustomMetadataIsNull() { RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); @@ -846,8 +836,9 @@ public void WriteWithoutCustomMetadata_LastBatchCustomMetadataIsNull() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - reader.ReadNextRecordBatch(); - Assert.Null(reader.LastBatchCustomMetadata); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + Assert.Null(read.CustomMetadata); } [Fact] @@ -868,12 +859,16 @@ public void WriteCustomMetadata_MixedBatches_WithAndWithoutMetadata() using var reader = new ArrowStreamReader(stream); - reader.ReadNextRecordBatch(); - Assert.NotNull(reader.LastBatchCustomMetadata); - Assert.Equal("value", reader.LastBatchCustomMetadata["key"]); + Assert.Equal(meta, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); + + RecordBatchWithMetadata second = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(second.Batch); + Assert.Null(second.CustomMetadata); - reader.ReadNextRecordBatch(); - Assert.Null(reader.LastBatchCustomMetadata); + // At the end of the stream both halves are null, not the previous batch's metadata. + RecordBatchWithMetadata end = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.Null(end.Batch); + Assert.Null(end.CustomMetadata); } [Fact] @@ -891,8 +886,9 @@ public void WriteCustomMetadata_EmptyDictionary_WritesNoMetadata() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - reader.ReadNextRecordBatch(); - Assert.Null(reader.LastBatchCustomMetadata); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + Assert.Null(read.CustomMetadata); } [Fact] @@ -953,10 +949,10 @@ public void WriteCustomMetadata_RejectedMetadata_LeavesWriterUsable() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - RecordBatch readBatch = reader.ReadNextRecordBatch(); - Assert.NotNull(readBatch); - ArrowReaderVerifier.CompareBatches(batch, readBatch); - Assert.Equal(good, reader.LastBatchCustomMetadata); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + ArrowReaderVerifier.CompareBatches(batch, read.Batch); + Assert.Equal(good, read.CustomMetadata); Assert.Null(reader.ReadNextRecordBatch()); } @@ -976,9 +972,10 @@ public void WriteCustomMetadata_EmptyValues_RoundTrips() stream.Position = 0; using var reader = new ArrowStreamReader(stream); - reader.ReadNextRecordBatch(); - Assert.NotNull(reader.LastBatchCustomMetadata); - Assert.Equal("", reader.LastBatchCustomMetadata["empty"]); + IReadOnlyDictionary readMetadata = + reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata; + Assert.NotNull(readMetadata); + Assert.Equal("", readMetadata["empty"]); } /// diff --git a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs index 22317231..58cbbd71 100644 --- a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs +++ b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs @@ -119,11 +119,11 @@ public void ImportCustomMetadata_PythonWrites() using var ms = new MemoryStream(ipcBytes); using var reader = new ArrowStreamReader(ms); - RecordBatch batch = reader.ReadNextRecordBatch(); + (RecordBatch batch, IReadOnlyDictionary metadata) = + reader.ReadNextRecordBatchWithCustomMetadata(); Assert.NotNull(batch); Assert.Equal(5, batch.Length); - var metadata = reader.LastBatchCustomMetadata; Assert.NotNull(metadata); Assert.Equal("python", metadata["origin"]); Assert.Equal("2", metadata["version"]);