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.