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). ///