From c450bdfb6fab5c5697edc9def7543784e8192040 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Fri, 16 Oct 2026 11:27:45 -0400 Subject: [PATCH] feat(ipc): Add ArrowIpcWriterWriteArrayViewWithDictionaries A writer that produces batches one at a time, rather than from an ArrowArrayStream, had no public way to write the DictionaryBatch messages a dictionary-encoded batch depends on: ArrowIpcWriterWriteArrayStream() did this internally, but only for a whole stream. Expose that per-batch path, which ArrowIpcWriterWriteArrayStream() now uses as well. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 17 +++- src/nanoarrow/ipc/writer_test.cc | 165 ++++++++++++++++++++++++++++++- src/nanoarrow/nanoarrow_ipc.h | 18 ++++ 3 files changed, 193 insertions(+), 7 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index bb8c79710..1f0c8f246 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -408,6 +408,16 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( return NANOARROW_OK; } +ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( + struct ArrowIpcWriter* writer, const struct ArrowArrayView* in, + struct ArrowError* error) { + NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && in != NULL); + int64_t next_dictionary_id = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( + writer, in, &next_dictionary_id, error)); + return ArrowIpcWriterWriteArrayView(writer, in, error); +} + static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl( struct ArrowIpcWriter* writer, struct ArrowArrayStream* in, struct ArrowSchema* schema, struct ArrowArray* array, @@ -424,11 +434,8 @@ static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl( NANOARROW_RETURN_NOT_OK(ArrowArrayViewSetArray(array_view, array, error)); - int64_t next_dictionary_id = 0; - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( - writer, array_view, &next_dictionary_id, error)); - - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteArrayView(writer, array_view, error)); + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterWriteArrayViewWithDictionaries(writer, array_view, error)); ArrowArrayRelease(array); } diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index edcc8839c..538e908c8 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -303,7 +303,8 @@ TEST(NanoarrowIpcWriter, WriteDictionaryBatch) { // Build a struct array with a single dictionary-encoded (int32 -> utf8) child. static void MakeDictionaryStructArray(struct ArrowArray* array, - struct ArrowSchema* schema) { + struct ArrowSchema* schema, + const char* value1 = "bar") { 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), @@ -320,7 +321,7 @@ static void MakeDictionaryStructArray(struct ArrowArray* array, ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("foo")), NANOARROW_OK); - ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("bar")), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView(value1)), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendInt(indices, 1), NANOARROW_OK); @@ -726,3 +727,163 @@ TEST(NanoarrowIpcWriter, CompressedDictionaryStream) { EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo"); EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar"); } + +// The type of each message in an encapsulated IPC stream, up to the end of the stream +static std::vector ReadMessageTypes( + const struct ArrowBuffer* stream_bytes) { + std::vector types; + 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; + } + types.push_back(decoder->message_type); + int64_t advance = decoder->header_size_bytes + decoder->body_size_bytes; + view.data.as_uint8 += advance; + view.size_bytes -= advance; + } + return types; +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesMatchesArrayStream) { + // Writing a batch with ArrowIpcWriterWriteArrayViewWithDictionaries() produces the + // same bytes as writing it through ArrowIpcWriterWriteArrayStream() + struct ArrowError error; + + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(array.get(), schema.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 by_view; + { + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), by_view.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; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK); + } + + nanoarrow::UniqueBuffer by_stream; + { + nanoarrow::UniqueSchema stream_schema; + ASSERT_EQ(ArrowSchemaDeepCopy(schema.get(), stream_schema.get()), NANOARROW_OK); + nanoarrow::UniqueArrayStream array_stream; + ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), stream_schema.get(), 1), + NANOARROW_OK); + ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get()); + + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), by_stream.get()), + NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error), + NANOARROW_OK) + << error.message; + } + + EXPECT_EQ( + ReadMessageTypes(by_view.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + ASSERT_EQ(by_view->size_bytes, by_stream->size_bytes); + EXPECT_EQ(memcmp(by_view->data, by_stream->data, by_view->size_bytes), 0); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) { + // Batches written one at a time each carry their own dictionary, which the reader + // resolves against the DictionaryBatch written just before it + struct ArrowError error; + + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray first; + ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(first.get(), schema.get(), "bar")); + nanoarrow::UniqueSchema second_schema; + nanoarrow::UniqueArray second; + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(second.get(), second_schema.get(), "baz")); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.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; + for (struct ArrowArray* array : {first.get(), second.get()}) { + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + } + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK); + + EXPECT_EQ( + ReadMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + + struct ArrowIpcInputStream input; + ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, output.get()), NANOARROW_OK); + nanoarrow::UniqueArrayStream reader; + ASSERT_EQ(ArrowIpcArrayStreamReaderInit(reader.get(), &input, nullptr), NANOARROW_OK); + nanoarrow::UniqueSchema roundtrip_schema; + ASSERT_EQ(ArrowArrayStreamGetSchema(reader.get(), roundtrip_schema.get(), &error), + NANOARROW_OK) + << error.message; + + for (const char* expected : {"bar", "baz"}) { + SCOPED_TRACE(expected); + nanoarrow::UniqueArray roundtrip; + ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_NE(roundtrip->release, nullptr); + + nanoarrow::UniqueArrayView view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), roundtrip_schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(view.get(), roundtrip.get(), &error), NANOARROW_OK) + << error.message; + struct ArrowArrayView* values = view->children[0]->dictionary; + ASSERT_NE(values, nullptr); + struct ArrowStringView value1 = ArrowArrayViewGetStringUnsafe(values, 1); + EXPECT_EQ(std::string(value1.data, value1.size_bytes), expected); + } +} diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 7f5559759..dc1526ce8 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -130,6 +130,8 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteSchema) #define ArrowIpcWriterWriteArrayView \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayView) +#define ArrowIpcWriterWriteArrayViewWithDictionaries \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayViewWithDictionaries) #define ArrowIpcWriterWriteDictionaryBatch \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteDictionaryBatch) #define ArrowIpcWriterWriteArrayStream \ @@ -1163,6 +1165,22 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* const struct ArrowArrayView* in, struct ArrowError* error); +/// \brief Write an array view and the dictionaries it references +/// +/// Writes a DictionaryBatch message for each dictionary-encoded field of in, followed +/// by the RecordBatch itself, exactly as ArrowIpcWriterWriteArrayStream() does for each +/// array of a stream. This lets a caller that produces batches one at a time (rather +/// than from an ArrowArrayStream) write dictionary-encoded data. Dictionary IDs are +/// assigned in the same order as when the schema was written. Each dictionary is +/// written in full (not as a delta) before every RecordBatch, and file mode has the +/// same limitations as ArrowIpcWriterWriteArrayStream(). in must not be NULL; use +/// ArrowIpcWriterWriteArrayView() to write EOS. +/// +/// Errors are propagated from the underlying encoder and output byte stream. +NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( + struct ArrowIpcWriter* writer, const struct ArrowArrayView* in, + struct ArrowError* error); + /// \brief Write a DictionaryBatch message to the output byte stream /// /// dictionary_id must match the id assigned to the dictionary-encoded field in the