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 284d3003ceb3580a25b6a0ac5357f813b0fdcd32 Author: Curt Hagenlocher <[email protected]> AuthorDate: Sun Sep 6 21:06:42 2026 -0700 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) <[email protected]> Claude-Session: https://claude.ai/code/session_01XhMo3XSWYHo1apHd9PzZTb --- .../Internal/FlightDataStream.cs | 4 +- src/Apache.Arrow/Ipc/ArrowFileWriter.cs | 21 ---- src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs | 19 ++-- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 82 ++++++++-------- test/Apache.Arrow.Tests/ArrowFileWriterTests.cs | 108 +++++++++++++++++++++ test/Apache.Arrow.Tests/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 3ab4331..46a7ac2 100644 --- a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs +++ b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs @@ -55,7 +55,7 @@ namespace Apache.Arrow.Flight.Internal 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 @@ namespace Apache.Arrow.Flight.Internal _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 91b7c29..cfdf226 100644 --- a/src/Apache.Arrow/Ipc/ArrowFileWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowFileWriter.cs @@ -66,27 +66,6 @@ namespace Apache.Arrow.Ipc RecordBatchBlocks = new List<Block>(); } - 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 acac52f..45fd792 100644 --- a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs +++ b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs @@ -166,19 +166,16 @@ namespace Apache.Arrow.Ipc private static IReadOnlyDictionary<string, string> ReadMessageCustomMetadata(Flatbuf.Message message) { - int count = message.CustomMetadataLength; - if (count == 0) - return null; - - var result = new Dictionary<string, string>(count); - for (int i = 0; i < count; i++) + Dictionary<string, string> metadata = message.CustomMetadataLength > 0 + ? new Dictionary<string, string>(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<byte> buffer) diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index c8e6032..aa63985 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -805,14 +805,17 @@ namespace Apache.Arrow.Ipc Builder, compressionType, Flatbuf.BodyCompressionMethod.BUFFER); } - 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 + // 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 @@ 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) - { - ValidateCustomMetadata(customMetadata); - Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata); - customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); - } + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); // Serialize record batch @@ -861,16 +857,17 @@ namespace Apache.Arrow.Ipc 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) { + // 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 @@ 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) - { - ValidateCustomMetadata(customMetadata); - Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata); - customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); - } + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); // Serialize record batch @@ -1090,7 +1080,7 @@ namespace Apache.Arrow.Ipc 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 @@ namespace Apache.Arrow.Ipc public virtual void WriteRecordBatch(RecordBatch recordBatch) { - WriteRecordBatchInternal(recordBatch); + WriteRecordBatchInternal(recordBatch, customMetadata: null); } public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata) @@ -1170,7 +1160,7 @@ namespace Apache.Arrow.Ipc 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<string, string> customMetadata, CancellationToken cancellationToken = default) @@ -1332,12 +1322,32 @@ namespace Apache.Arrow.Ipc return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered); } + /// <summary> + /// Builds the Message-level custom_metadata vector, or a default offset when there is none. + /// </summary> + private VectorOffset GetCustomMetadataOffset(IReadOnlyDictionary<string, string> customMetadata) + { + if (customMetadata == null || customMetadata.Count == 0) + { + return default; + } + + Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata); + return Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets); + } + /// <summary> /// 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. /// </summary> private static void ValidateCustomMetadata(IReadOnlyDictionary<string, string> customMetadata) { + if (customMetadata == null) + { + return; + } + foreach (KeyValuePair<string, string> metadatum in customMetadata) { if (metadatum.Key == null) @@ -1394,7 +1404,7 @@ namespace Apache.Arrow.Ipc // 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 @@ 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, diff --git a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs index d810a53..af5eee8 100644 --- a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs @@ -311,6 +311,114 @@ namespace Apache.Arrow.Tests 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<string, string> { ["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<string, string> { ["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<string, string> + { + ["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<string, string> { ["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<string, string> { ["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 1df27fd..5246e1a 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 @@ namespace Apache.Arrow.Tests 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<string, string>()); + 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<string, string> rejects a null key, so go through a map that allows one. + var withNullKey = new NullTolerantMetadata(new KeyValuePair<string, string>(null, "value")); + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + Assert.Throws<ArgumentException>(() => writer.WriteRecordBatch(batch, withNullKey)); + } + + [Fact] + public void WriteCustomMetadata_NullValue_Throws() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary<string, string> { ["key"] = null }; + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + Assert.Throws<ArgumentException>(() => writer.WriteRecordBatch(batch, meta)); + } + + [Fact] + public async Task WriteCustomMetadataAsync_NullValue_Throws() + { + RecordBatch batch = TestData.CreateSampleRecordBatch(length: 5); + var meta = new Dictionary<string, string> { ["key"] = null }; + + using var stream = new MemoryStream(); + using var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true); + + await Assert.ThrowsAsync<ArgumentException>( + () => 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<string, string> { ["key"] = "value" }; + var bad = new Dictionary<string, string> { ["key"] = null }; + + using var stream = new MemoryStream(); + using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true)) + { + Assert.Throws<ArgumentException>(() => 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 @@ namespace Apache.Arrow.Tests Assert.NotNull(reader.LastBatchCustomMetadata); Assert.Equal("", reader.LastBatchCustomMetadata["empty"]); } + + /// <summary> + /// A metadata collection that can hold a null key, which <see cref="Dictionary{TKey, TValue}"/> cannot. + /// </summary> + private sealed class NullTolerantMetadata : IReadOnlyDictionary<string, string> + { + private readonly KeyValuePair<string, string>[] _entries; + + public NullTolerantMetadata(params KeyValuePair<string, string>[] entries) => _entries = entries; + + public int Count => _entries.Length; + public IEnumerable<string> Keys => _entries.Select(e => e.Key); + public IEnumerable<string> 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<KeyValuePair<string, string>> GetEnumerator() => ((IEnumerable<KeyValuePair<string, string>>)_entries).GetEnumerator(); + IEnumerator IEnumerable.GetEnumerator() => _entries.GetEnumerator(); + } } }
