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 304d018488c029586191447b17031377959514a6 Author: Curt Hagenlocher <[email protected]> AuthorDate: Mon Sep 7 13:25:37 2026 -0700 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<RecordBatchWithMetadata> 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) <[email protected]> Claude-Session: https://claude.ai/code/session_01XhMo3XSWYHo1apHd9PzZTb --- src/Apache.Arrow/Ipc/ArrowFileReader.cs | 11 ++++ src/Apache.Arrow/Ipc/ArrowStreamReader.cs | 37 +++++++++-- src/Apache.Arrow/Ipc/RecordBatchWithMetadata.cs | 49 ++++++++++++++ test/Apache.Arrow.Tests/ArrowFileWriterTests.cs | 15 +++-- test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs | 77 +++++++++++----------- .../CustomMetadataPythonTests.cs | 4 +- 6 files changed, 141 insertions(+), 52 deletions(-) diff --git a/src/Apache.Arrow/Ipc/ArrowFileReader.cs b/src/Apache.Arrow/Ipc/ArrowFileReader.cs index fa3f84a..c8b63ef 100644 --- a/src/Apache.Arrow/Ipc/ArrowFileReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowFileReader.cs @@ -85,5 +85,16 @@ namespace Apache.Arrow.Ipc { return Implementation.ReadRecordBatchAsync(index, cancellationToken); } + + /// <summary> + /// 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. + /// </summary> + public async ValueTask<RecordBatchWithMetadata> 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 e7b5861..bdc1fb7 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamReader.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamReader.cs @@ -154,12 +154,37 @@ namespace Apache.Arrow.Ipc } /// <summary> - /// 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 <see cref="ArrowStreamWriter.WriteRecordBatch(RecordBatch, IReadOnlyDictionary{string, string})"/>. /// </summary> - public IReadOnlyDictionary<string, string> LastBatchCustomMetadata => _implementation.LastBatchCustomMetadata; + /// <returns> + /// The record batch and its custom metadata. At the end of the stream both + /// <see cref="RecordBatchWithMetadata.Batch"/> and + /// <see cref="RecordBatchWithMetadata.CustomMetadata"/> are null; the metadata is also + /// null for a batch whose message carried none. + /// </returns> + public async ValueTask<RecordBatchWithMetadata> ReadNextRecordBatchWithCustomMetadataAsync(CancellationToken cancellationToken = default) + { + RecordBatch batch = await _implementation.ReadNextRecordBatchAsync(cancellationToken).ConfigureAwait(false); + + return batch == null ? default : new RecordBatchWithMetadata(batch, _implementation.LastBatchCustomMetadata); + } + + /// <summary> + /// Reads the next record batch together with the custom metadata on its IPC Message, + /// the counterpart of <see cref="ArrowStreamWriter.WriteRecordBatch(RecordBatch, IReadOnlyDictionary{string, string})"/>. + /// </summary> + /// <returns> + /// The record batch and its custom metadata. At the end of the stream both + /// <see cref="RecordBatchWithMetadata.Batch"/> and + /// <see cref="RecordBatchWithMetadata.CustomMetadata"/> are null; the metadata is also + /// null for a batch whose message carried none. + /// </returns> + 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 0000000..54e2103 --- /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 +{ + /// <summary> + /// A record batch read from an Arrow IPC source, together with the custom metadata + /// carried on the IPC Message that held it. + /// </summary> + public readonly struct RecordBatchWithMetadata + { + public RecordBatchWithMetadata(RecordBatch batch, IReadOnlyDictionary<string, string> customMetadata) + { + Batch = batch; + CustomMetadata = customMetadata; + } + + /// <summary> + /// The record batch that was read, or null at the end of the stream. + /// </summary> + public RecordBatch Batch { get; } + + /// <summary> + /// The Message-level custom metadata accompanying <see cref="Batch"/>, or null if the + /// message carried none. + /// </summary> + public IReadOnlyDictionary<string, string> CustomMetadata { get; } + + public void Deconstruct(out RecordBatch batch, out IReadOnlyDictionary<string, string> customMetadata) + { + batch = Batch; + customMetadata = CustomMetadata; + } + } +} diff --git a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs index af5eee8..f3b0a34 100644 --- a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs @@ -372,8 +372,14 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests 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 5246e1a..172f69b 100644 --- a/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs @@ -759,16 +759,15 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests stream.Position = 0; using var reader = new ArrowStreamReader(stream); - RecordBatch readBatch = reader.ReadNextRecordBatch(); + (RecordBatch readBatch, IReadOnlyDictionary<string, string> 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 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests 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 @@ namespace Apache.Arrow.Tests stream.Position = 0; using var reader = new ArrowStreamReader(stream); - reader.ReadNextRecordBatch(); - Assert.NotNull(reader.LastBatchCustomMetadata); - Assert.Equal("", reader.LastBatchCustomMetadata["empty"]); + IReadOnlyDictionary<string, string> readMetadata = + reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata; + Assert.NotNull(readMetadata); + Assert.Equal("", readMetadata["empty"]); } /// <summary> diff --git a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs index 2231723..58cbbd7 100644 --- a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs +++ b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs @@ -119,11 +119,11 @@ namespace Apache.Arrow.Tests using var ms = new MemoryStream(ipcBytes); using var reader = new ArrowStreamReader(ms); - RecordBatch batch = reader.ReadNextRecordBatch(); + (RecordBatch batch, IReadOnlyDictionary<string, string> 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"]);
