feat: Support IPC Message custom_metadata in ArrowStreamWriter and ArrowStreamReader

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<long> WriteMessageAsync<T>(MessageHeader headerType, Offset<T> headerOffset, int bodyLength, CancellationToken cancellationToken)
+        private protected override ValueTask<long> WriteMessageAsync<T>(MessageHeader headerType, Offset<T> headerOffset, int bodyLength, VectorOffset customMetadataOffset, CancellationToken cancellationToken)
         {
             Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage(
                 Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
-                bodyLength);
+                bodyLength, customMetadataOffset);
 
             Builder.Finish(messageOffset.Value);
 

src/Apache.Arrow/Ipc/ArrowFileReader.cs

@@ -85,5 +85,16 @@ public ValueTask<RecordBatch> ReadRecordBatchAsync(int index, CancellationToken
         {
             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);
+        }
     }
 }

src/Apache.Arrow/Ipc/ArrowFileWriter.cs

@@ -66,27 +66,6 @@ public ArrowFileWriter(Stream stream, Schema schema, bool leaveOpen, IpcOptions
             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;

src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs

@@ -81,6 +81,11 @@ protected virtual void Dispose(bool disposing)
         public abstract ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken);
         public abstract RecordBatch ReadNextRecordBatch();
 
+        /// <summary>
+        /// Custom metadata from the most recently read RecordBatch Message, if any.
+        /// </summary>
+        internal IReadOnlyDictionary<string, string> LastBatchCustomMetadata { get; private protected set; }
+
         internal static T ReadMessage<T>(ByteBuffer bb)
             where T : struct, IFlatbufferObject
         {
@@ -148,6 +153,7 @@ protected RecordBatch CreateArrowObjectFromMessage(
                     }
 
                     List<IArrowArray> 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<string, string> ReadMessageCustomMetadata(Flatbuf.Message message)
+        {
+            Dictionary<string, string> metadata = message.CustomMetadataLength > 0
+                ? new Dictionary<string, string>(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<byte> buffer)
         {
             return new ByteBuffer(new ReadOnlyMemoryBufferAllocator(buffer), 0);

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();
         }
+
+        /// <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 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);
+        }
     }
 }

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<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)
             {
@@ -829,6 +837,8 @@ private protected void WriteRecordBatchInternal(RecordBatch recordBatch)
 
             VectorOffset buffersVectorOffset = Builder.EndVector();
 
+            VectorOffset customMetadataVectorOffset = GetCustomMetadataOffset(customMetadata);
+
             // Serialize record batch
 
             StartingWritingRecordBatch();
@@ -840,16 +850,24 @@ 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 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);
@@ -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<string, string> 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<string, string> 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);
         }
 
+        /// <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 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)
+                {
+                    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<Flatbuf.KeyValue>[] GetMetadataOffsets(IReadOnlyDictionary<string, string> 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.
         /// </returns>
         private protected long WriteMessage<T>(
-            Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength)
+            Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength,
+            VectorOffset customMetadataOffset = default)
             where T : struct
         {
             Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage(
                 Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
-                bodyLength);
+                bodyLength, customMetadataOffset);
 
             Builder.Finish(messageOffset.Value);
 
@@ -1378,12 +1449,13 @@ private protected long WriteMessage<T>(
         /// </returns>
         private protected virtual async ValueTask<long> WriteMessageAsync<T>(
             Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength,
+            VectorOffset customMetadataOffset,
             CancellationToken cancellationToken)
             where T : struct
         {
             Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage(
                 Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
-                bodyLength);
+                bodyLength, customMetadataOffset);
 
             Builder.Finish(messageOffset.Value);