You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by pa...@apache.org on 2026/10/20 15:34:41 UTC
(arrow-nanoarrow) branch main updated: feat: Add ArrowIpcWriterSetMessageMetadata (#967)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new d01fdbb5 feat: Add ArrowIpcWriterSetMessageMetadata (#967)
d01fdbb5 is described below
commit d01fdbb50cfdf838afcaed20307e58127d63deac
Author: Rusty Conover <ru...@conover.me>
AuthorDate: Tue Oct 20 11:34:35 2026 -0400
feat: Add ArrowIpcWriterSetMessageMetadata (#967)
`ArrowIpcEncoderSetMessageMetadata()` (#924, narrowed in #937) attaches
`Message::custom_metadata` to the next message an encoder encodes. An
`ArrowIpcWriter` keeps its encoder private, though, so a caller using
the writer has no way to attach metadata to the RecordBatch messages it
writes. The only option is to drop down to the encoder and redo the
writer's framing, footer bookkeeping and dictionary handling by hand.
This PR adds `ArrowIpcWriterSetMessageMetadata(writer, metadata,
error)`. The writer holds the metadata and attaches it to the next
RecordBatch message written by `ArrowIpcWriterWriteArrayView()` or
`ArrowIpcWriterWriteArrayStream()`. The metadata applies to that one
message. Passing `NULL` clears metadata that has been set but not yet
written. Like the encoder function, it takes ownership of the buffer.
The metadata is held by the writer rather than set on the encoder right
away. Setting it on the encoder would attach it to whichever message is
encoded next, and that can be a DictionaryBatch written ahead of the
batch. This happens both with an explicit
`ArrowIpcWriterWriteDictionaryBatch()` call and inside
`ArrowIpcWriterWriteArrayStream()`. RecordBatch-level metadata belongs
on the RecordBatch, so the writer passes it to the encoder just before
encoding that batch.
Things worth a look in review:
- With `WriteArrayStream()` the metadata lands on the first RecordBatch
only. Per-batch metadata over a stream would need a callback or a
per-array side channel, which seemed out of scope. The batch-at-a-time
pattern is `SetMessageMetadata()` + `WriteArrayView()`.
- #928 makes encoder-level metadata apply to DictionaryBatch messages
too. The two compose: this PR never sets encoder metadata until right
before the RecordBatch is encoded. They touch neighbouring code in
`writer.c`, so whichever lands second will need a small rebase. I can do
that.
Motivation: I'm moving the IPC layer of a DuckDB extension from Arrow
C++ to nanoarrow. Its RPC protocol carries per-batch state in
RecordBatch `custom_metadata` (what Arrow C++ writes for
`RecordBatchWriter::WriteRecordBatch(batch, custom_metadata)`), and this
was the one thing the writer couldn't express.
Tests:
- `MessageMetadataAttachesToNextRecordBatch` sets metadata, writes a
DictionaryBatch and then a RecordBatch, and checks that only the
RecordBatch carries it. It also checks that the following batch doesn't
inherit the metadata, and that set-then-cleared metadata is never
written.
- `MessageMetadataWithArrayStream` covers `WriteArrayStream()` writing
its own dictionaries.
ctest passes locally (macOS arm64, Debug, ZSTD/LZ4, 390 tests). The IPC
writer, encoder and files tests are clean under ASan/UBSan.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5.5 (1M context) <no...@anthropic.com>
---
src/nanoarrow/ipc/writer.c | 32 +++++++
src/nanoarrow/ipc/writer_test.cc | 193 +++++++++++++++++++++++++++++++++++++++
src/nanoarrow/nanoarrow_ipc.h | 17 ++++
3 files changed, 242 insertions(+)
diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c
index 1f0c8f24..9879f1b6 100644
--- a/src/nanoarrow/ipc/writer.c
+++ b/src/nanoarrow/ipc/writer.c
@@ -181,6 +181,11 @@ struct ArrowIpcWriterPrivate {
int writing_file;
int64_t bytes_written;
struct ArrowIpcFooter footer;
+
+ // Message custom_metadata held for the next RecordBatch message (not for any
+ // DictionaryBatch messages written ahead of it)
+ struct ArrowBuffer pending_message_metadata;
+ int has_pending_message_metadata;
};
ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer,
@@ -202,6 +207,8 @@ ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer,
private->writing_file = 0;
private->bytes_written = 0;
ArrowIpcFooterInit(&private->footer);
+ ArrowBufferInit(&private->pending_message_metadata);
+ private->has_pending_message_metadata = 0;
writer->private_data = private;
return NANOARROW_OK;
@@ -220,6 +227,7 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) {
ArrowBufferReset(&private->body_buffer);
ArrowIpcFooterReset(&private->footer);
+ ArrowBufferReset(&private->pending_message_metadata);
ArrowFree(private);
}
@@ -237,6 +245,24 @@ ArrowErrorCode ArrowIpcWriterSetCompression(struct ArrowIpcWriter* writer,
compression_level, error);
}
+ArrowErrorCode ArrowIpcWriterSetMessageMetadata(struct ArrowIpcWriter* writer,
+ struct ArrowBuffer* metadata,
+ struct ArrowError* error) {
+ NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL);
+ NANOARROW_UNUSED(error);
+ struct ArrowIpcWriterPrivate* private =
+ (struct ArrowIpcWriterPrivate*)writer->private_data;
+
+ ArrowBufferReset(&private->pending_message_metadata);
+ private->has_pending_message_metadata = 0;
+ if (metadata != NULL) {
+ ArrowBufferMove(metadata, &private->pending_message_metadata);
+ private->has_pending_message_metadata = 1;
+ }
+
+ return NANOARROW_OK;
+}
+
static struct ArrowBufferView ArrowBufferToBufferView(const struct ArrowBuffer* buffer) {
struct ArrowBufferView buffer_view = {
.data.as_uint8 = buffer->data,
@@ -300,6 +326,12 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* writer,
NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0));
NANOARROW_ASSERT_OK(ArrowBufferResize(&private->body_buffer, 0, 0));
+ if (private->has_pending_message_metadata) {
+ private->has_pending_message_metadata = 0;
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderSetMessageMetadata(
+ &private->encoder, &private->pending_message_metadata, error));
+ }
+
NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSimpleRecordBatch(
&private->encoder, in, &private->body_buffer, error));
NANOARROW_RETURN_NOT_OK_WITH_ERROR(
diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc
index 538e908c..c01e0998 100644
--- a/src/nanoarrow/ipc/writer_test.cc
+++ b/src/nanoarrow/ipc/writer_test.cc
@@ -19,6 +19,10 @@
#include <stdio.h>
+#include <string>
+#include <utility>
+#include <vector>
+
#include "nanoarrow/nanoarrow_ipc.hpp"
TEST(NanoarrowIpcWriter, OutputStreamBuffer) {
@@ -887,3 +891,192 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) {
EXPECT_EQ(std::string(value1.data, value1.size_bytes), expected);
}
}
+
+namespace {
+
+// Walk the messages of an encapsulated IPC stream, recording each message's type and
+// custom_metadata (packed like ArrowSchema::metadata; empty if the message has none)
+std::vector<std::pair<enum ArrowIpcMessageType, std::string>> ReadMessageMetadata(
+ struct ArrowBuffer* stream_bytes) {
+ std::vector<std::pair<enum ArrowIpcMessageType, std::string>> out;
+ nanoarrow::ipc::UniqueDecoder decoder;
+ EXPECT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);
+
+ struct ArrowError error;
+ struct ArrowBufferView view;
+ view.data.as_uint8 = stream_bytes->data;
+ view.size_bytes = stream_bytes->size_bytes;
+
+ while (view.size_bytes > 0) {
+ ArrowErrorCode result = ArrowIpcDecoderVerifyHeader(decoder.get(), view, &error);
+ if (result == ENODATA) {
+ break;
+ }
+ EXPECT_EQ(result, NANOARROW_OK) << error.message;
+ if (result != NANOARROW_OK) {
+ break;
+ }
+
+ nanoarrow::UniqueBuffer metadata;
+ EXPECT_EQ(ArrowIpcDecoderGetMessageMetadata(decoder.get(), metadata.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ out.emplace_back(decoder->message_type,
+ std::string(reinterpret_cast<const char*>(metadata->data),
+ static_cast<size_t>(metadata->size_bytes)));
+
+ int64_t advance = decoder->header_size_bytes + decoder->body_size_bytes;
+ view.data.as_uint8 += advance;
+ view.size_bytes -= advance;
+ }
+
+ return out;
+}
+
+// metadata with a single "state" key, returning its packed form
+std::string MakeStateMetadata(struct ArrowBuffer* metadata, const char* value) {
+ EXPECT_EQ(ArrowMetadataBuilderInit(metadata, nullptr), NANOARROW_OK);
+ EXPECT_EQ(
+ ArrowMetadataBuilderAppend(metadata, ArrowCharView("state"), ArrowCharView(value)),
+ NANOARROW_OK);
+ return std::string(reinterpret_cast<const char*>(metadata->data),
+ static_cast<size_t>(metadata->size_bytes));
+}
+
+// struct<d: dictionary<int32, string>> with two rows
+void MakeDictionaryBatch(struct ArrowSchema* schema, struct ArrowArray* array) {
+ struct ArrowError error;
+ ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "d"), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, &error), NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(array->children[0]->dictionary, ArrowCharView("foo")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(array->children[0]->dictionary, ArrowCharView("bar")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishElement(array), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 0), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishElement(array), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, &error), NANOARROW_OK)
+ << error.message;
+}
+
+} // namespace
+
+TEST(NanoarrowIpcWriter, MessageMetadataAttachesToNextRecordBatch) {
+ struct ArrowError error;
+
+ nanoarrow::UniqueSchema schema;
+ nanoarrow::UniqueArray array;
+ ASSERT_NO_FATAL_FAILURE(MakeDictionaryBatch(schema.get(), array.get()));
+
+ nanoarrow::UniqueArrayView array_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK);
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), NANOARROW_OK)
+ << error.message;
+
+ // Metadata set before a DictionaryBatch is held for the RecordBatch that follows it
+ nanoarrow::UniqueBuffer metadata;
+ std::string first = MakeStateMetadata(metadata.get(), "first");
+ ASSERT_EQ(ArrowIpcWriterSetMessageMetadata(writer.get(), metadata.get(), &error),
+ NANOARROW_OK);
+ EXPECT_EQ(metadata->data, nullptr);
+ ASSERT_EQ(ArrowIpcWriterWriteDictionaryBatch(
+ writer.get(), /*dictionary_id=*/0,
+ /*is_delta=*/0, array_view->children[0]->dictionary, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ // The next RecordBatch does not inherit it
+ ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ // Metadata that is set and then cleared is never written
+ nanoarrow::UniqueBuffer cleared;
+ MakeStateMetadata(cleared.get(), "cleared");
+ ASSERT_EQ(ArrowIpcWriterSetMessageMetadata(writer.get(), cleared.get(), &error),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterSetMessageMetadata(writer.get(), nullptr, &error),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK);
+
+ auto messages = ReadMessageMetadata(output.get());
+ ASSERT_EQ(messages.size(), 5);
+ EXPECT_EQ(messages[0].first, NANOARROW_IPC_MESSAGE_TYPE_SCHEMA);
+ EXPECT_EQ(messages[0].second, "");
+ EXPECT_EQ(messages[1].first, NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH);
+ EXPECT_EQ(messages[1].second, "");
+ EXPECT_EQ(messages[2].first, NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH);
+ EXPECT_EQ(messages[2].second, first);
+ EXPECT_EQ(messages[3].first, NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH);
+ EXPECT_EQ(messages[3].second, "");
+ EXPECT_EQ(messages[4].first, NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH);
+ EXPECT_EQ(messages[4].second, "");
+}
+
+TEST(NanoarrowIpcWriter, MessageMetadataWithArrayStream) {
+ // ArrowIpcWriterWriteArrayStream() writes the stream's dictionaries itself; metadata
+ // set beforehand lands on the first RecordBatch only
+ struct ArrowError error;
+
+ nanoarrow::UniqueSchema schema;
+ nanoarrow::UniqueArray array;
+ ASSERT_NO_FATAL_FAILURE(MakeDictionaryBatch(schema.get(), array.get()));
+
+ nanoarrow::UniqueArrayStream array_stream;
+ nanoarrow::UniqueSchema stream_schema;
+ ASSERT_EQ(ArrowSchemaDeepCopy(schema.get(), stream_schema.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), stream_schema.get(), 1),
+ NANOARROW_OK);
+ ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get());
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK);
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
+
+ nanoarrow::UniqueBuffer metadata;
+ std::string first = MakeStateMetadata(metadata.get(), "first");
+ ASSERT_EQ(ArrowIpcWriterSetMessageMetadata(writer.get(), metadata.get(), &error),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ auto messages = ReadMessageMetadata(output.get());
+ ASSERT_EQ(messages.size(), 3);
+ EXPECT_EQ(messages[0].first, NANOARROW_IPC_MESSAGE_TYPE_SCHEMA);
+ EXPECT_EQ(messages[0].second, "");
+ EXPECT_EQ(messages[1].first, NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH);
+ EXPECT_EQ(messages[1].second, "");
+ EXPECT_EQ(messages[2].first, NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH);
+ EXPECT_EQ(messages[2].second, first);
+}
diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h
index bc287fd6..a607df80 100644
--- a/src/nanoarrow/nanoarrow_ipc.h
+++ b/src/nanoarrow/nanoarrow_ipc.h
@@ -126,6 +126,8 @@
#define ArrowIpcWriterReset NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterReset)
#define ArrowIpcWriterSetCompression \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterSetCompression)
+#define ArrowIpcWriterSetMessageMetadata \
+ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterSetMessageMetadata)
#define ArrowIpcWriterWriteSchema \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteSchema)
#define ArrowIpcWriterWriteArrayView \
@@ -1156,6 +1158,21 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterSetCompression(
struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type,
int compression_level, struct ArrowError* error);
+/// \brief Set the custom metadata of the next RecordBatch message
+///
+/// The metadata (Message::custom_metadata, in the representation built by
+/// ArrowMetadataBuilderInit()/ArrowMetadataBuilderAppend()) is held by the writer and
+/// attached to the next RecordBatch message written by ArrowIpcWriterWriteArrayView()
+/// or ArrowIpcWriterWriteArrayStream(). It is not attached to any DictionaryBatch
+/// message written before that RecordBatch, so it can be set ahead of a call that
+/// writes dictionaries and a batch together. The writer takes ownership of metadata
+/// (it is moved and left empty); pass NULL to clear metadata that has been set but not
+/// yet written. If a write returns an error, metadata that was set may still be held
+/// for a later RecordBatch; pass NULL to clear it if that matters.
+NANOARROW_DLL ArrowErrorCode
+ArrowIpcWriterSetMessageMetadata(struct ArrowIpcWriter* writer,
+ struct ArrowBuffer* metadata, struct ArrowError* error);
+
/// \brief Write a schema to the output byte stream
///
/// Errors are propagated from the underlying encoder and output byte stream.