From aafc9b08b23717ac58e79af1da3d46d6b170fb77 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Fri, 16 Oct 2026 11:04:14 -0400 Subject: [PATCH] feat(ipc): Add ArrowIpcWriterSetMessageMetadata Message::custom_metadata can be attached through the encoder (ArrowIpcEncoderSetMessageMetadata), but an ArrowIpcWriter's encoder is private, so a writer could not attach metadata to the RecordBatch messages it writes. The new function holds the metadata in the writer and attaches it to the next RecordBatch, not to DictionaryBatch messages written ahead of it (by ArrowIpcWriterWriteDictionaryBatch() or within ArrowIpcWriterWriteArrayStream()). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 32 +++++ src/nanoarrow/ipc/writer_test.cc | 193 +++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 16 +++ 3 files changed, 241 insertions(+) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index bb8c79710..3a8915682 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 edcc8839c..12b3544ff 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -19,6 +19,10 @@ #include +#include +#include +#include + #include "nanoarrow/nanoarrow_ipc.hpp" TEST(NanoarrowIpcWriter, OutputStreamBuffer) { @@ -726,3 +730,192 @@ TEST(NanoarrowIpcWriter, CompressedDictionaryStream) { EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo"); EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar"); } + +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> ReadMessageMetadata( + struct ArrowBuffer* stream_bytes) { + std::vector> 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(metadata->data), + static_cast(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(metadata->data), + static_cast(metadata->size_bytes)); +} + +// struct> 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 7f5559759..e189904d9 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 \ @@ -1146,6 +1148,20 @@ 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. +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.