diff --git a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs index 50b2a40b..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; @@ -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/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/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 2c380e5d..45fd7920 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,20 @@ protected RecordBatch CreateArrowObjectFromMessage( return null; } + private static IReadOnlyDictionary ReadMessageCustomMetadata(Flatbuf.Message message) + { + Dictionary metadata = message.CustomMetadataLength > 0 + ? new Dictionary(message.CustomMetadataLength) : null; + for (int i = 0; i < message.CustomMetadataLength; i++) + { + Flatbuf.KeyValue keyValue = message.CustomMetadata(i).GetValueOrDefault(); + + metadata[keyValue.Key] = keyValue.Value; + } + + return metadata; + } + 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..bdc1fb7f 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,39 @@ public RecordBatch ReadNextRecordBatch() { return _implementation.ReadNextRecordBatch(); } + + /// + /// 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 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/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index a39caa66..aa63985e 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -805,9 +805,17 @@ public ArrowStreamWriter(Stream baseStream, Schema schema, bool leaveOpen, IpcOp Builder, compressionType, Flatbuf.BodyCompressionMethod.BUFFER); } - private protected void WriteRecordBatchInternal(RecordBatch recordBatch) + 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) { @@ -829,6 +837,8 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch) VectorOffset buffersVectorOffset = Builder.EndVector(); + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); + // Serialize record batch StartingWritingRecordBatch(); @@ -840,7 +850,7 @@ 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); @@ -848,8 +858,16 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch) } 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); @@ -870,6 +888,8 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat VectorOffset buffersVectorOffset = Builder.EndVector(); + VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata); + // Serialize record batch StartingWritingRecordBatch(); @@ -882,6 +902,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); @@ -1059,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); @@ -1129,12 +1150,22 @@ 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) + { + WriteRecordBatchInternal(recordBatch, customMetadata); } 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) + { + return WriteRecordBatchInternalAsync(recordBatch, customMetadata, cancellationToken); } public void WriteStart() @@ -1291,6 +1322,45 @@ 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 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) + { + 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); @@ -1334,7 +1404,7 @@ private VectorOffset GetFieldMetadataOffset(Field field) // Build message - await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, cancellationToken) + await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, default, cancellationToken) .ConfigureAwait(false); return schemaOffset; @@ -1347,12 +1417,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); @@ -1378,12 +1449,13 @@ private protected long WriteMessage( /// 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/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 d810a53b..f3b0a343 100644 --- a/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs +++ b/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs @@ -311,6 +311,121 @@ 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); + 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] + 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); + RecordBatchWithMetadata read = await reader.ReadNextRecordBatchWithCustomMetadataAsync(); + Assert.NotNull(read.Batch); + Assert.Equal(customMetadata, read.CustomMetadata); + } + + [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 1a4b5a6b..172f69b7 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; @@ -736,5 +737,264 @@ 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); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + ArrowReaderVerifier.CompareBatches(originalBatch, read.Batch); + + 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] + 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, IReadOnlyDictionary readMetadata) = + await reader.ReadNextRecordBatchWithCustomMetadataAsync(); + Assert.NotNull(readBatch); + ArrowReaderVerifier.CompareBatches(originalBatch, readBatch); + + Assert.Equal(customMetadata, readMetadata); + } + + [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); + + Assert.Equal(meta1, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); + Assert.Equal(meta2, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); + } + + [Fact] + public void WriteWithoutCustomMetadata_CustomMetadataIsNull() + { + 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); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + Assert.Null(read.CustomMetadata); + } + + [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); + + Assert.Equal(meta, reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata); + + RecordBatchWithMetadata second = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(second.Batch); + Assert.Null(second.CustomMetadata); + + // 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] + 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); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + Assert.Null(read.CustomMetadata); + } + + [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); + RecordBatchWithMetadata read = reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(read.Batch); + ArrowReaderVerifier.CompareBatches(batch, read.Batch); + Assert.Equal(good, read.CustomMetadata); + Assert.Null(reader.ReadNextRecordBatch()); + } + + [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); + IReadOnlyDictionary readMetadata = + reader.ReadNextRecordBatchWithCustomMetadata().CustomMetadata; + Assert.NotNull(readMetadata); + Assert.Equal("", readMetadata["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(); + } } } diff --git a/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs new file mode 100644 index 00000000..58cbbd71 --- /dev/null +++ b/test/Apache.Arrow.Tests/CustomMetadataPythonTests.cs @@ -0,0 +1,132 @@ +// 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 + // ------------------------------------------------------------------- + + [Collection("PythonNet")] + public class CustomMetadataPythonTests + { + public CustomMetadataPythonTests(PythonNetFixture pythonNet) + { + pythonNet.EnsureInitialized(); + } + + // ------------------------------------------------------------------- + // 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, IReadOnlyDictionary metadata) = + reader.ReadNextRecordBatchWithCustomMetadata(); + Assert.NotNull(batch); + Assert.Equal(5, batch.Length); + + Assert.NotNull(metadata); + Assert.Equal("python", metadata["origin"]); + Assert.Equal("2", metadata["version"]); + } + } +}