Skip to content

Support IPC Message custom_metadata in ArrowStreamWriter and ArrowStreamReader (rebased continuation of #283) #424

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rustyconover wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Support IPC Message custom_metadata in ArrowStreamWriter and ArrowStreamReader (rebased continuation of #283) #424

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a123f03
Expose IPC Message custom_metadata on ArrowStreamReader
Mar 8, 2026
9260dfe
Add WriteRecordBatch overload with custom_metadata
Mar 8, 2026
fcb77e5
Add custom_metadata cross-language tests using pythonnet + pyarrow:
Mar 11, 2026
0177e8b
Address Copilot review comments on #424
rustyconover Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/Apache.Arrow.Flight/Internal/FlightDataStream.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
23 changes: 23 additions & 0 deletions src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down Expand Up @@ -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
Expand All @@ -158,6 +164,23 @@ protected RecordBatch CreateArrowObjectFromMessage(
return null;
}

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++)
{
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<byte> buffer)
{
return new ByteBuffer(new ReadOnlyMemoryBufferAllocator(buffer), 0);
Expand Down
10 changes: 10 additions & 0 deletions src/Apache.Arrow/Ipc/ArrowStreamReader.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
// limitations under the License.

using System;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
Expand Down Expand Up @@ -151,5 +152,14 @@ public RecordBatch ReadNextRecordBatch()
{
return _implementation.ReadNextRecordBatch();
}

/// <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.
/// </summary>
public IReadOnlyDictionary<string, string> LastBatchCustomMetadata => _implementation.LastBatchCustomMetadata;
}
}
78 changes: 74 additions & 4 deletions src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, string> customMetadata)
{
// TODO: Truncate buffers with extraneous padding / unused capacity

Expand All @@ -829,6 +834,15 @@ 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)
{
ValidateCustomMetadata(customMetadata);
Offset<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata);
customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets);
Comment on lines +837 to +843

@rustyconover rustyconover Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 0177e8b: added a ValidateCustomMetadata check that throws a clear ArgumentException for null keys/values before building the FlatBuffer offsets.

}

// Serialize record batch

StartingWritingRecordBatch();
Expand All @@ -840,14 +854,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,

@CurtHagenlocher CurtHagenlocher Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My suspicion is that it's better not to have this overload. We don't need it for backwards-compatibility and removing it might avoid an error in a derived class.

CancellationToken cancellationToken = default)
{
return WriteRecordBatchInternalAsync(recordBatch, customMetadata: null, cancellationToken);
}

private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBatch,
IReadOnlyDictionary<string, string> customMetadata,
CancellationToken cancellationToken = default)
{
if (!HasWrittenSchema)
Expand All @@ -870,6 +891,15 @@ 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<Flatbuf.KeyValue>[] metadataOffsets = GetMetadataOffsets(customMetadata);
customMetadataVectorOffset = Flatbuf.Message.CreateCustomMetadataVector(Builder, metadataOffsets);
}

// Serialize record batch

StartingWritingRecordBatch();
Expand All @@ -882,6 +912,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);
Expand Down Expand Up @@ -1132,11 +1163,21 @@ public virtual void WriteRecordBatch(RecordBatch recordBatch)
WriteRecordBatchInternal(recordBatch);
}

public virtual void WriteRecordBatch(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata)
{
WriteRecordBatchInternal(recordBatch, customMetadata);

@CurtHagenlocher CurtHagenlocher Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Based on my previous analysis of the original PR, this overload breaks ArrowFileWriter because it would skip the call to WriteStart which is performed in that class's WriteRecordBatch/WriteRecordBatchAsync overrides.

}

public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default)
{
return WriteRecordBatchInternalAsync(recordBatch, cancellationToken);
}

public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, IReadOnlyDictionary<string, string> customMetadata, CancellationToken cancellationToken = default)
{
return WriteRecordBatchInternalAsync(recordBatch, customMetadata, cancellationToken);
}

public void WriteStart()
{
if (!HasWrittenStart)
Expand Down Expand Up @@ -1291,6 +1332,25 @@ private VectorOffset GetFieldMetadataOffset(Field field)
return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered);
}

/// <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.
/// </summary>
private static void ValidateCustomMetadata(IReadOnlyDictionary<string, string> customMetadata)
{
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);
Expand Down Expand Up @@ -1347,12 +1407,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);

Expand All @@ -1376,14 +1437,23 @@ private protected long WriteMessage<T>(
/// <returns>
/// The number of bytes written to the stream.
/// </returns>
private protected virtual ValueTask<long> WriteMessageAsync<T>(

@CurtHagenlocher CurtHagenlocher Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider making this non-virtual or removing it entirely and forcing use of the signature with the customMetadataOffset (which is what WriteMessage does).

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,
CancellationToken cancellationToken)
where T : struct
{
Offset<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage(
Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
bodyLength);
bodyLength, customMetadataOffset);

Builder.Finish(messageOffset.Value);

Expand Down
Loading