feat(ipc): Add ArrowIpcWriterSetMessageMetadata

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(

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) {
@@ -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<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);
+}

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.