From 8f8e296cf022587c1d926d930db1a21a61a8bd12 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 19 Oct 2026 15:25:21 -0400 Subject: [PATCH 1/2] feat(ipc): Skip unchanged dictionaries in ArrowIpcWriterWriteArrayViewWithDictionaries ArrowIpcWriterWriteArrayViewWithDictionaries() wrote every dictionary before every RecordBatch: its input is borrowed, so it couldn't compare a dictionary with the previous batch's by address as ArrowIpcWriterWriteArrayStream() does. The writer now keeps a copy of each dictionary it writes and compares the next batch's values with it, writing only dictionaries that changed (and those whose values contain a nested dictionary that was written). In a file, batches whose dictionaries are unchanged can now be written this way. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 313 ++++++++++++++++++++++++++++++- src/nanoarrow/ipc/writer_test.cc | 301 ++++++++++++++++++++++++++++- src/nanoarrow/nanoarrow_ipc.h | 21 ++- 3 files changed, 618 insertions(+), 17 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 18252125c..f570af58a 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -190,8 +190,51 @@ struct ArrowIpcWriterPrivate { // DictionaryBatch messages written ahead of it) struct ArrowBuffer pending_message_metadata; int has_pending_message_metadata; + // Copies of the dictionaries last written by + // ArrowIpcWriterWriteArrayViewWithDictionaries() (struct ArrowIpcWriterViewDictionary) + struct ArrowBuffer view_dictionaries; }; +// The values of a dictionary as last written by +// ArrowIpcWriterWriteArrayViewWithDictionaries(). Its input is borrowed, so the +// values are copied to compare them with the next batch's. +struct ArrowIpcWriterViewDictionary { + int64_t dictionary_id; + struct ArrowArray values; +}; + +static void ArrowIpcWriterForgetViewDictionary(struct ArrowIpcWriterPrivate* private, + int64_t dictionary_id) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + if (entries[i].dictionary_id == dictionary_id) { + ArrowArrayRelease(&entries[i].values); + entries[i] = entries[n_entries - 1]; + NANOARROW_ASSERT_OK( + ArrowBufferResize(&private->view_dictionaries, + private->view_dictionaries.size_bytes - + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary), + /*shrink_to_fit=*/0)); + return; + } + } +} + +static void ArrowIpcWriterResetViewDictionaries(struct ArrowIpcWriterPrivate* private) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + ArrowArrayRelease(&entries[i].values); + } + NANOARROW_ASSERT_OK( + ArrowBufferResize(&private->view_dictionaries, 0, /*shrink_to_fit=*/0)); +} + ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, struct ArrowIpcOutputStream* output_stream) { NANOARROW_DCHECK(writer != NULL && output_stream != NULL); @@ -214,6 +257,7 @@ ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, ArrowBufferInit(&private->file_dictionary_ids); ArrowBufferInit(&private->pending_message_metadata); private->has_pending_message_metadata = 0; + ArrowBufferInit(&private->view_dictionaries); writer->private_data = private; return NANOARROW_OK; @@ -234,6 +278,8 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { ArrowIpcFooterReset(&private->footer); ArrowBufferReset(&private->file_dictionary_ids); ArrowBufferReset(&private->pending_message_metadata); + ArrowIpcWriterResetViewDictionaries(private); + ArrowBufferReset(&private->view_dictionaries); ArrowFree(private); } @@ -300,6 +346,8 @@ ArrowErrorCode ArrowIpcWriterWriteSchema(struct ArrowIpcWriter* writer, // A dictionary may be replaced in a stream (e.g., when the dictionaries of // ArrowIpcWriterWriteArrayStream() input change between batches) but not in a file ArrowIpcEncoderSetDictionaryReplacement(&private->encoder, !private->writing_file); + // Dictionary IDs are assigned again for the new schema + ArrowIpcWriterResetViewDictionaries(private); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0)); NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSchema(&private->encoder, in, error)); @@ -373,10 +421,9 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* writer, return NANOARROW_OK; } -ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( +static ArrowErrorCode ArrowIpcWriterWriteDictionaryBatchImpl( struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta, const struct ArrowArrayView* values_view, struct ArrowError* error) { - NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && values_view != NULL); struct ArrowIpcWriterPrivate* private = (struct ArrowIpcWriterPrivate*)writer->private_data; @@ -433,6 +480,20 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( return NANOARROW_OK; } +ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( + struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta, + const struct ArrowArrayView* values_view, struct ArrowError* error) { + NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && values_view != NULL); + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + + // ArrowIpcWriterWriteArrayViewWithDictionaries() must not skip this ID's next + // dictionary because it matches the values it wrote before this one + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + return ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, is_delta, + values_view, error); +} + // Check whether two arrays are guaranteed to contain identical values because they // have the same shape and point to the same buffers (e.g., a dictionary shared by // several arrays in a stream). This is only meaningful while both arrays are alive: @@ -517,15 +578,255 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArray( return NANOARROW_OK; } +// Initialize out as an array view with the same types as src +static ArrowErrorCode ArrowIpcWriterArrayViewInitLike(struct ArrowArrayView* out, + const struct ArrowArrayView* src) { + ArrowArrayViewInitFromType(out, src->storage_type); + out->layout = src->layout; + NANOARROW_RETURN_NOT_OK(ArrowArrayViewAllocateChildren(out, src->n_children)); + for (int64_t i = 0; i < src->n_children; i++) { + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterArrayViewInitLike(out->children[i], src->children[i])); + } + + if (src->dictionary != NULL) { + NANOARROW_RETURN_NOT_OK(ArrowArrayViewAllocateDictionary(out)); + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterArrayViewInitLike(out->dictionary, src->dictionary)); + } + + return NANOARROW_OK; +} + +// Producers need not initialize the padding bits of a bitmap's last byte. Clear them +// so that they don't make two copies of the same values compare unequal. +static void ArrowIpcWriterClearBitmapPadding(struct ArrowArray* array, + const struct ArrowArrayView* array_view) { + int64_t remainder = array->length % 8; + if (remainder != 0) { + uint8_t mask = (uint8_t)((1U << remainder) - 1U); + for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) { + if (array_view->layout.element_size_bits[i] == 1) { + struct ArrowBuffer* buffer = ArrowArrayBuffer(array, i); + if (buffer->size_bytes > 0) { + buffer->data[buffer->size_bytes - 1] &= mask; + } + } + } + } + + for (int64_t i = 0; i < array->n_children; i++) { + ArrowIpcWriterClearBitmapPadding(array->children[i], array_view->children[i]); + } +} + +// Copy the values of a dictionary into out, in a form that only depends on the values: +// zero offset, buffers sized to the values, and cleared bitmap padding. A dictionary +// nested in the values is not copied: it is written (and compared) separately. +static ArrowErrorCode ArrowIpcWriterCopyDictionaryValues( + const struct ArrowArrayView* values_view, struct ArrowArray* out, + struct ArrowError* error) { + NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromArrayView(out, values_view, error)); + ArrowErrorCode result = ArrowArrayStartAppending(out); + if (result == NANOARROW_OK) { + result = ArrowArrayReserve(out, values_view->length); + } + if (result == NANOARROW_OK) { + result = ArrowArrayAppendStorageFromArrayView(out, values_view, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayFinishBuildingDefault(out, error); + } + if (result != NANOARROW_OK) { + ArrowArrayRelease(out); + return result; + } + + ArrowIpcWriterClearBitmapPadding(out, values_view); + return NANOARROW_OK; +} + +// Compare two copies made by ArrowIpcWriterCopyDictionaryValues() from views of +// the same type as values_view +static ArrowErrorCode ArrowIpcWriterDictionaryValuesEqual( + const struct ArrowArrayView* values_view, const struct ArrowArray* lhs, + const struct ArrowArray* rhs, int* out, struct ArrowError* error) { + struct ArrowArrayView lhs_view; + struct ArrowArrayView rhs_view; + ArrowArrayViewInitFromType(&lhs_view, NANOARROW_TYPE_UNINITIALIZED); + ArrowArrayViewInitFromType(&rhs_view, NANOARROW_TYPE_UNINITIALIZED); + + ArrowErrorCode result = ArrowIpcWriterArrayViewInitLike(&lhs_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowIpcWriterArrayViewInitLike(&rhs_view, values_view); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&lhs_view, lhs, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&rhs_view, rhs, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewCompare(&lhs_view, &rhs_view, NANOARROW_COMPARE_IDENTICAL, out, + NULL); + } + + ArrowArrayViewReset(&lhs_view); + ArrowArrayViewReset(&rhs_view); + return result; +} + +static struct ArrowIpcWriterViewDictionary* ArrowIpcWriterFindViewDictionary( + struct ArrowIpcWriterPrivate* private, int64_t dictionary_id) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + if (entries[i].dictionary_id == dictionary_id) { + return &entries[i]; + } + } + return NULL; +} + +// Check whether values_view is physically identical to copy (made by +// ArrowIpcWriterCopyDictionaryValues() from a view of the same type). If it is, the +// values are equal; if not, they may still be (e.g., if values_view has an offset). +static ArrowErrorCode ArrowIpcWriterViewMatchesCopy( + const struct ArrowArrayView* values_view, const struct ArrowArray* copy, int* out, + struct ArrowError* error) { + struct ArrowArrayView copy_view; + ArrowArrayViewInitFromType(©_view, NANOARROW_TYPE_UNINITIALIZED); + ArrowErrorCode result = ArrowIpcWriterArrayViewInitLike(©_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(©_view, copy, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewCompare(values_view, ©_view, NANOARROW_COMPARE_IDENTICAL, + out, NULL); + } + ArrowArrayViewReset(©_view); + return result; +} + +// Write the dictionary values_view for dictionary_id unless its values are identical +// to the ones this function last wrote for that ID (and force is not set). Sets +// *written to whether a DictionaryBatch was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( + struct ArrowIpcWriter* writer, int64_t dictionary_id, + const struct ArrowArrayView* values_view, int force, int* written, + struct ArrowError* error) { + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + + struct ArrowIpcWriterViewDictionary* entry = + ArrowIpcWriterFindViewDictionary(private, dictionary_id); + + // The common case: the values are physically identical to the copy (e.g., because + // the producer builds each batch's dictionary the same way) + if (entry != NULL && !force) { + int identical = 0; + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterViewMatchesCopy(values_view, &entry->values, &identical, error)); + if (identical) { + *written = 0; + return NANOARROW_OK; + } + } + + // Otherwise, compare a copy of the values made the same way (zero offset, cleared + // padding bits, and computed null count) + struct ArrowArray values; + if (ArrowIpcWriterCopyDictionaryValues(values_view, &values, error) != NANOARROW_OK) { + // These values can't be copied (e.g., a dense union), so they can't be compared + // with the next batch's either: write them every time + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + *written = 1; + return ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/0, values_view, error); + } + + ArrowErrorCode result = NANOARROW_OK; + if (entry != NULL && !force) { + int equal = 0; + result = ArrowIpcWriterDictionaryValuesEqual(values_view, &entry->values, &values, + &equal, error); + if (result == NANOARROW_OK && equal) { + ArrowArrayRelease(&values); + *written = 0; + return NANOARROW_OK; + } + } + + if (result == NANOARROW_OK) { + result = ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/0, values_view, error); + } + + if (result == NANOARROW_OK && entry == NULL) { + struct ArrowIpcWriterViewDictionary new_entry = {.dictionary_id = dictionary_id}; + ArrowArrayMove(&values, &new_entry.values); + result = + ArrowBufferAppend(&private->view_dictionaries, &new_entry, sizeof(new_entry)); + if (result != NANOARROW_OK) { + ArrowArrayRelease(&new_entry.values); + } + } else if (result == NANOARROW_OK) { + ArrowArrayRelease(&entry->values); + ArrowArrayMove(&values, &entry->values); + } else { + ArrowArrayRelease(&values); + } + + *written = result == NANOARROW_OK; + return result; +} + +// Write the dictionaries of a borrowed array view in the same order as +// ArrowIpcWriterWriteDictionariesForArray(), skipping those that are identical to +// the ones last written for their IDs. A dictionary whose values contain a dictionary +// that was written is written again even if its own values are unchanged: a reader +// resolves a DictionaryBatch against the nested dictionaries current at that point. +// Sets *written if any dictionary of array_view was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionaries( + struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view, + int64_t* next_id, int* written, struct ArrowError* error) { + int64_t dictionary_id = *next_id; + if (array_view->dictionary != NULL) { + (*next_id)++; + } + + for (int64_t i = 0; i < array_view->n_children; i++) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, array_view->children[i], next_id, written, error)); + } + + if (array_view->dictionary != NULL) { + int nested_written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, array_view->dictionary, next_id, &nested_written, error)); + + int dictionary_written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionary( + writer, dictionary_id, array_view->dictionary, /*force=*/nested_written, + &dictionary_written, error)); + if (nested_written || dictionary_written) { + *written = 1; + } + } + + 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; - // in is borrowed, so its dictionaries can't be kept alive to compare against the - // next batch's: write every one of them - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( - writer, in, /*array=*/NULL, /*previous=*/NULL, &next_dictionary_id, error)); + int written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, in, &next_dictionary_id, &written, error)); return ArrowIpcWriterWriteArrayView(writer, in, error); } diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index a67cfb718..10077ad92 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -1480,9 +1480,9 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) { } TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesInFile) { - // A dictionary ID may only be written once in a file: writing one batch with its - // dictionaries succeeds, and a second is rejected rather than writing an invalid - // file + // A dictionary ID may only be written once in a file: batches whose dictionaries + // are unchanged can be written, and a changed dictionary is rejected rather than + // writing an invalid file struct ArrowError error; nanoarrow::UniqueSchema schema; @@ -1507,9 +1507,304 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesInFile) { &error), NANOARROW_OK) << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + + nanoarrow::UniqueSchema changed_schema; + nanoarrow::UniqueArray changed; + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(changed.get(), changed_schema.get(), "baz")); + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), changed.get(), &error), NANOARROW_OK) + << error.message; EXPECT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), &error), ENOTSUP); + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK) + << error.message; + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchFileReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + EXPECT_EQ(maybe_arrow_reader.ValueUnsafe()->num_record_batches(), 2); +#endif +} + +// Write each array with ArrowIpcWriterWriteArrayViewWithDictionaries() to a stream +static void WriteArrayViews(struct ArrowSchema* schema, + const std::vector& arrays, + struct ArrowBuffer* output) { + struct ArrowError error; + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema, &error), NANOARROW_OK) + << error.message; + + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema, &error), NANOARROW_OK) + << error.message; + for (struct ArrowArray* array : arrays) { + 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); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesSkipsUnchanged) { + // Separately built arrays with the same dictionary values don't share any buffers, + // but only a changed dictionary is written again + std::vector schemas(4); + std::vector arrays(4); + const char* values[] = {"bar", "bar", "baz", "baz"}; + for (size_t i = 0; i < arrays.size(); i++) { + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(arrays[i].get(), schemas[i].get(), values[i])); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews( + schemas[0].get(), + {arrays[0].get(), arrays[1].get(), arrays[2].get(), arrays[3].get()}, + output.get())); + EXPECT_EQ(DecodeMessageTypes(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_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchStreamReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + for (const char* expected : values) { + std::shared_ptr arrow_batch; + ASSERT_TRUE(maybe_arrow_reader.ValueUnsafe()->ReadNext(&arrow_batch).ok()); + ASSERT_NE(arrow_batch, nullptr); + auto arrow_dictionary = + std::static_pointer_cast(arrow_batch->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + EXPECT_EQ(arrow_values->GetString(1), expected); + } +#endif + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 4); + for (size_t i = 0; i < roundtrip.size(); i++) { + EXPECT_EQ(StringValues(roundtrip[i]->children[0]->dictionary), + (std::vector{"foo", values[i]})); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesComparesValues) { + // Dictionaries are compared by value: a slice of a longer dictionary, an unknown null + // count, and garbage in a bitmap's padding bits don't make an unchanged dictionary + // look changed + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ASSERT_EQ( + ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_INT64), + NANOARROW_OK); + + // Values [1, null, 3], built as is (first) or as a slice of [0, 1, null, 3, 4] + std::vector arrays(2); + for (size_t i = 0; i < arrays.size(); i++) { + struct ArrowArray* array = arrays[i].get(); + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema.get(), nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendInt(values, 0), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 3), NANOARROW_OK); + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendInt(values, 4), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 2), NANOARROW_OK); + array->length = 2; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, &error), NANOARROW_OK) + << error.message; + if (i == 1) { + // Set the padding bits of the validity bitmap, slice to [1, null, 3], and leave + // the null count unknown + ArrowArrayBuffer(values, 0)->data[0] |= 0xe0; + values->offset = 1; + values->length = 3; + values->null_count = -1; + } + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schema.get(), {arrays[0].get(), arrays[1].get()}, output.get())); + EXPECT_EQ(DecodeMessageTypes(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_RECORD_BATCH})); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNested) { + // A changed nested dictionary is written, and so is the dictionary whose values + // refer to it, even though that dictionary's own values are unchanged + const std::vector inner_values = {"bar", "bar", "baz"}; + std::vector schemas(inner_values.size()); + std::vector arrays(inner_values.size()); + for (size_t i = 0; i < inner_values.size(); i++) { + ASSERT_NO_FATAL_FAILURE(MakeNestedDictionaryStructArray( + arrays[i].get(), schemas[i].get(), inner_values[i].c_str())); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schemas[0].get(), + {arrays[0].get(), arrays[1].get(), arrays[2].get()}, output.get())); + EXPECT_EQ(DecodeDictionaryIds(output.get()), (std::vector{1, 0, 1, 0})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 3); + for (size_t i = 0; i < roundtrip.size(); i++) { + struct ArrowArray* outer_dictionary = roundtrip[i]->children[0]->dictionary; + EXPECT_EQ(StringValues(outer_dictionary->children[0]->dictionary), + (std::vector{"foo", inner_values[i]})); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesAfterWriteDictionaryBatch) { + // A dictionary written with ArrowIpcWriterWriteDictionaryBatch() replaces the one + // ArrowIpcWriterWriteArrayViewWithDictionaries() wrote, so the next batch's + // dictionary is written even though it is identical to that one + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(array.get(), schema.get())); + nanoarrow::UniqueSchema other_schema; + nanoarrow::UniqueArray other; + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(other.get(), other_schema.get(), "baz")); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView other_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(other_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(other_view.get(), other.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; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcWriterWriteDictionaryBatch(writer.get(), 0, /*is_delta=*/0, + other_view->children[0]->dictionary, &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(DecodeMessageTypes(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_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); + EXPECT_EQ(StringValues(roundtrip[1]->children[0]->dictionary), + (std::vector{"foo", "bar"})); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesUncopyableValues) { + // Dense union values can't be copied to compare them with the next batch's, so they + // are written for every batch + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ArrowSchemaInit(schema->children[0]->dictionary); + ASSERT_EQ(ArrowSchemaSetTypeUnion(schema->children[0]->dictionary, + NANOARROW_TYPE_DENSE_UNION, 1), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema->children[0]->dictionary->children[0], + NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0]->dictionary->children[0], "i"), + NANOARROW_OK); + + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + for (int32_t value : {5, 6}) { + ASSERT_EQ(ArrowArrayAppendInt(values->children[0], value), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishUnionElement(values, 0), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 1), NANOARROW_OK); + array->length = 1; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), &error), NANOARROW_OK) + << error.message; + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schema.get(), {array.get(), array.get()}, output.get())); + EXPECT_EQ(DecodeMessageTypes(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})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); } namespace { diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index d10f2fc18..1fe9b3620 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1211,12 +1211,16 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// assigned in the same order as when the schema was written, and nested dictionaries /// are written before the dictionaries that refer to them. /// -/// Because in is borrowed, its dictionaries are not compared with those of earlier -/// batches (as ArrowIpcWriterWriteArrayStream() does to skip unchanged ones): each -/// dictionary is written in full before every RecordBatch, replacing the previous -/// one. In file mode, where a dictionary ID may only be written once, this means only -/// one batch with dictionaries can be written this way. in must not be NULL; use -/// ArrowIpcWriterWriteArrayView() to write EOS. +/// A dictionary whose values are identical to the ones this function last wrote for +/// its ID is not written again. Because in is borrowed, the writer keeps a copy of +/// each dictionary it writes to compare with the next batch's (so the comparison costs +/// a copy of each dictionary per batch); values that can't be copied (e.g., a dense +/// union) are written every time. A changed dictionary is written in full, replacing +/// the previous one, which a file does not support (ENOTSUP). A dictionary whose +/// values contain a dictionary that was written is written again too. +/// ArrowIpcWriterWriteDictionaryBatch() and ArrowIpcWriterWriteSchema() clear the +/// copies they affect. 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( @@ -1228,8 +1232,9 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( /// dictionary_id must match the id assigned to the dictionary-encoded field in the /// schema. is_delta selects DictionaryBatch.isDelta. values_view must not itself be /// dictionary-encoded. The writer does not check that a schema was already written. -/// values_view is always written: the writer only sees a borrowed view and cannot -/// check whether it is identical to a dictionary that was already written. When +/// values_view is always written: unlike ArrowIpcWriterWriteArrayViewWithDictionaries(), +/// this does not check whether it is identical to a dictionary that was already +/// written. When /// writing a file, returns ENOTSUP for a delta or for a dictionary_id that was already /// written (IPC files do not support dictionary replacement). /// From 6bc6f0cdbb598ea87ce213ceb01a521c0a52697d Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 19 Oct 2026 15:37:19 -0400 Subject: [PATCH 2/2] feat(ipc): Write dictionary deltas from ArrowIpcWriterWriteArrayViewWithDictionaries When a batch's dictionary starts with the values last written for its ID, the writer now writes only the appended values, as a delta DictionaryBatch, instead of replacing the whole dictionary. This only applies to streams: a file does not support deltas. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 113 ++++++++++++++-- src/nanoarrow/ipc/writer_test.cc | 215 +++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 8 +- 3 files changed, 323 insertions(+), 13 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index f570af58a..4866756dd 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -620,19 +620,25 @@ static void ArrowIpcWriterClearBitmapPadding(struct ArrowArray* array, } } -// Copy the values of a dictionary into out, in a form that only depends on the values: -// zero offset, buffers sized to the values, and cleared bitmap padding. A dictionary -// nested in the values is not copied: it is written (and compared) separately. +// Copy length values of a dictionary, starting at offset, into out, in a form that +// only depends on the values: zero offset, buffers sized to the values, and cleared +// bitmap padding. A dictionary nested in the values is not copied: it is written (and +// compared) separately. static ArrowErrorCode ArrowIpcWriterCopyDictionaryValues( - const struct ArrowArrayView* values_view, struct ArrowArray* out, - struct ArrowError* error) { + const struct ArrowArrayView* values_view, int64_t offset, int64_t length, + struct ArrowArray* out, struct ArrowError* error) { + struct ArrowArrayView slice = *values_view; + slice.offset += offset; + slice.length = length; + slice.null_count = -1; + NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromArrayView(out, values_view, error)); ArrowErrorCode result = ArrowArrayStartAppending(out); if (result == NANOARROW_OK) { - result = ArrowArrayReserve(out, values_view->length); + result = ArrowArrayReserve(out, length); } if (result == NANOARROW_OK) { - result = ArrowArrayAppendStorageFromArrayView(out, values_view, error); + result = ArrowArrayAppendStorageFromArrayView(out, &slice, error); } if (result == NANOARROW_OK) { result = ArrowArrayFinishBuildingDefault(out, error); @@ -710,9 +716,84 @@ static ArrowErrorCode ArrowIpcWriterViewMatchesCopy( return result; } +// If values_view starts with the values last written for entry, write the values +// appended to them as a delta and add them to entry's copy. Sets *written to whether a +// delta was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionaryDelta( + struct ArrowIpcWriter* writer, struct ArrowIpcWriterViewDictionary* entry, + int64_t dictionary_id, const struct ArrowArrayView* values_view, int* written, + struct ArrowError* error) { + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + *written = 0; + + int64_t n_written = entry->values.length; + struct ArrowArray prefix; + if (ArrowIpcWriterCopyDictionaryValues(values_view, 0, n_written, &prefix, error) != + NANOARROW_OK) { + // Leave it to the caller, which writes values it can't copy in full + return NANOARROW_OK; + } + + int prefix_equal = 0; + ArrowErrorCode result = ArrowIpcWriterDictionaryValuesEqual( + values_view, &entry->values, &prefix, &prefix_equal, error); + if (result != NANOARROW_OK || !prefix_equal) { + ArrowArrayRelease(&prefix); + return result; + } + + struct ArrowArray delta; + result = ArrowIpcWriterCopyDictionaryValues( + values_view, n_written, values_view->length - n_written, &delta, error); + if (result != NANOARROW_OK) { + ArrowArrayRelease(&prefix); + return result; + } + + struct ArrowArrayView delta_view; + ArrowArrayViewInitFromType(&delta_view, NANOARROW_TYPE_UNINITIALIZED); + result = ArrowIpcWriterArrayViewInitLike(&delta_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&delta_view, &delta, error); + } + if (result == NANOARROW_OK) { + result = ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/1, &delta_view, error); + } + + if (result == NANOARROW_OK) { + *written = 1; + + // The values that were already written followed by the delta are a copy of + // values_view. If it can't be completed, forget the dictionary: the next batch's + // is then written in full. + ArrowErrorCode copy_result = + ArrowArrayAppendStorageFromArrayView(&prefix, &delta_view, error); + if (copy_result == NANOARROW_OK) { + copy_result = ArrowArrayFinishBuildingDefault(&prefix, error); + } + if (copy_result == NANOARROW_OK) { + ArrowIpcWriterClearBitmapPadding(&prefix, values_view); + ArrowArrayRelease(&entry->values); + ArrowArrayMove(&prefix, &entry->values); + } else { + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + } + } + + ArrowArrayViewReset(&delta_view); + ArrowArrayRelease(&delta); + if (prefix.release != NULL) { + ArrowArrayRelease(&prefix); + } + return result; +} + // Write the dictionary values_view for dictionary_id unless its values are identical -// to the ones this function last wrote for that ID (and force is not set). Sets -// *written to whether a DictionaryBatch was written. +// to the ones this function last wrote for that ID (and force is not set). In a stream, +// values that start with the ones last written are written as a delta of the values +// that were appended. Sets *written to whether a DictionaryBatch was written. static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( struct ArrowIpcWriter* writer, int64_t dictionary_id, const struct ArrowArrayView* values_view, int force, int* written, @@ -735,10 +816,22 @@ static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( } } + // In a stream, values that start with the ones last written only need the appended + // values to be written, as a delta (a file does not support deltas) + if (entry != NULL && !force && !private->writing_file && + values_view->length > entry->values.length) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaryDelta( + writer, entry, dictionary_id, values_view, written, error)); + if (*written) { + return NANOARROW_OK; + } + } + // Otherwise, compare a copy of the values made the same way (zero offset, cleared // padding bits, and computed null count) struct ArrowArray values; - if (ArrowIpcWriterCopyDictionaryValues(values_view, &values, error) != NANOARROW_OK) { + if (ArrowIpcWriterCopyDictionaryValues(values_view, 0, values_view->length, &values, + error) != NANOARROW_OK) { // These values can't be copied (e.g., a dense union), so they can't be compared // with the next batch's either: write them every time ArrowIpcWriterForgetViewDictionary(private, dictionary_id); diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 10077ad92..35fa9487e 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -1807,6 +1807,221 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesUncopyableValues) { ASSERT_EQ(roundtrip.size(), 2); } +// The isDelta flag of each DictionaryBatch message of an encapsulated IPC stream +static std::vector DecodeDictionaryDeltaFlags(const struct ArrowBuffer* buffer) { + std::vector is_delta; + struct ArrowBufferView remaining; + remaining.data.as_uint8 = buffer->data; + remaining.size_bytes = buffer->size_bytes; + nanoarrow::ipc::UniqueDecoder decoder; + EXPECT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowError error; + + while (remaining.size_bytes > 0) { + int result = ArrowIpcDecoderVerifyHeader(decoder.get(), remaining, &error); + if (result == ENODATA) { + break; + } + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + + if (decoder->message_type == NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH) { + result = ArrowIpcDecoderDecodeHeader(decoder.get(), remaining, &error); + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + is_delta.push_back(decoder->dictionary->is_delta); + } + + int64_t message_size = decoder->header_size_bytes + decoder->body_size_bytes; + remaining.data.as_uint8 += message_size; + remaining.size_bytes -= message_size; + } + return is_delta; +} + +// A struct> array of one row per dictionary value +static void MakeStringDictionaryStructArray(struct ArrowArray* array, + struct ArrowSchema* schema, + const std::vector& values) { + 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], "dict_col"), 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, nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + for (size_t i = 0; i < values.size(); i++) { + ASSERT_EQ(ArrowArrayAppendString(array->children[0]->dictionary, + ArrowCharView(values[i].c_str())), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], static_cast(i)), + NANOARROW_OK); + } + array->length = static_cast(values.size()); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesDeltas) { + // A dictionary that only grew is written as a delta of its new values; any other + // change replaces it + const std::vector> batch_values = { + {"foo", "bar"}, {"foo", "bar", "baz"}, {"foo", "bar", "baz"}, {"foo", "qux"}}; + std::vector schemas(batch_values.size()); + std::vector arrays(batch_values.size()); + std::vector array_ptrs; + for (size_t i = 0; i < arrays.size(); i++) { + ASSERT_NO_FATAL_FAILURE(MakeStringDictionaryStructArray( + arrays[i].get(), schemas[i].get(), batch_values[i])); + array_ptrs.push_back(arrays[i].get()); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews(schemas[0].get(), array_ptrs, output.get())); + EXPECT_EQ(DecodeMessageTypes(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, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + EXPECT_EQ(DecodeDictionaryDeltaFlags(output.get()), (std::vector{0, 1, 0})); + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchStreamReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + for (const auto& expected : batch_values) { + std::shared_ptr arrow_batch; + auto status = maybe_arrow_reader.ValueUnsafe()->ReadNext(&arrow_batch); + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(arrow_batch, nullptr); + auto arrow_dictionary = + std::static_pointer_cast(arrow_batch->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + ASSERT_EQ(arrow_values->length(), static_cast(expected.size())); + for (size_t i = 0; i < expected.size(); i++) { + EXPECT_EQ(arrow_values->GetString(static_cast(i)), expected[i]); + } + } +#endif + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), batch_values.size()); + for (size_t i = 0; i < roundtrip.size(); i++) { + EXPECT_EQ(StringValues(roundtrip[i]->children[0]->dictionary), batch_values[i]); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNullableDelta) { + // The appended values of a delta start partway through a byte of the validity + // bitmap, and the values that were already written still compare equal + struct ArrowError error; + std::vector schemas(2); + std::vector arrays(2); + for (size_t i = 0; i < arrays.size(); i++) { + struct ArrowSchema* schema = schemas[i].get(); + 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], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ASSERT_EQ( + ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_INT64), + NANOARROW_OK); + + struct ArrowArray* array = arrays[i].get(); + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + // [1, null, 3] and then [1, null, 3, null, 5] + ASSERT_EQ(ArrowArrayAppendInt(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 3), NANOARROW_OK); + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 5), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 0), NANOARROW_OK); + array->length = 1; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, &error), NANOARROW_OK) + << error.message; + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews( + schemas[0].get(), {arrays[0].get(), arrays[1].get()}, output.get())); + EXPECT_EQ(DecodeDictionaryDeltaFlags(output.get()), (std::vector{0, 1})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); + nanoarrow::UniqueArrayView view; + ArrowArrayViewInitFromType(view.get(), NANOARROW_TYPE_INT64); + ASSERT_EQ( + ArrowArrayViewSetArray(view.get(), roundtrip[1]->children[0]->dictionary, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(view->length, 5); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 0), 1); + EXPECT_TRUE(ArrowArrayViewIsNull(view.get(), 1)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 2), 3); + EXPECT_TRUE(ArrowArrayViewIsNull(view.get(), 3)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 4), 5); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNoDeltaInFile) { + // A file does not support deltas, so a dictionary that grew is a replacement there + struct ArrowError error; + std::vector schemas(2); + std::vector arrays(2); + ASSERT_NO_FATAL_FAILURE( + MakeStringDictionaryStructArray(arrays[0].get(), schemas[0].get(), {"foo", "bar"})); + ASSERT_NO_FATAL_FAILURE(MakeStringDictionaryStructArray( + arrays[1].get(), schemas[1].get(), {"foo", "bar", "baz"})); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schemas[0].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(ArrowIpcWriterStartFile(writer.get(), &error), NANOARROW_OK) << error.message; + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schemas[0].get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), arrays[0].get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), arrays[1].get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + ENOTSUP); +} + namespace { // Walk the messages of an encapsulated IPC stream, recording each message's type and diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 1fe9b3620..8c40daf62 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1215,9 +1215,11 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// its ID is not written again. Because in is borrowed, the writer keeps a copy of /// each dictionary it writes to compare with the next batch's (so the comparison costs /// a copy of each dictionary per batch); values that can't be copied (e.g., a dense -/// union) are written every time. A changed dictionary is written in full, replacing -/// the previous one, which a file does not support (ENOTSUP). A dictionary whose -/// values contain a dictionary that was written is written again too. +/// union) are written every time. In a stream, a dictionary that only grew (its values +/// start with the ones last written) is written as a delta of the appended values. +/// Otherwise, a changed dictionary is written in full, replacing the previous one, +/// which a file does not support (ENOTSUP). A dictionary whose +/// values contain a dictionary that was written is written again too, in full. /// ArrowIpcWriterWriteDictionaryBatch() and ArrowIpcWriterWriteSchema() clear the /// copies they affect. in must not be NULL; use ArrowIpcWriterWriteArrayView() to /// write EOS.