diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 2cccd33d0..ca70b244c 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -46,6 +46,7 @@ struct ArrowIpcEncoderPrivate { struct ArrowBuffer buffers; struct ArrowBuffer nodes; int encoding_footer; + int dictionary_replacement; struct ArrowIpcDictionaryEncodings dictionary_encodings; // Metadata to attach to the next encoded Message (in nanoarrow's packed // representation), or an empty buffer if the next Message has no metadata. @@ -76,6 +77,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { return ESPIPE; } private->encoding_footer = 0; + private->dictionary_replacement = 0; ArrowBufferInit(&private->buffers); ArrowBufferInit(&private->nodes); ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings); @@ -216,6 +218,14 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( return ArrowIpcEncoderSetCompressor(encoder, &compressor); } +void ArrowIpcEncoderSetDictionaryReplacement(struct ArrowIpcEncoder* encoder, + char enabled) { + NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL); + struct ArrowIpcEncoderPrivate* private = + (struct ArrowIpcEncoderPrivate*)encoder->private_data; + private->dictionary_replacement = enabled != 0; +} + static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out, size_t size) { _NANOARROW_CHECK_UPPER_LIMIT(size, INT32_MAX); @@ -554,6 +564,7 @@ static ArrowErrorCode ArrowIpcEncodeField( struct ArrowSchemaView schema_view; NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, schema, error)); + const struct ArrowSchema* value_schema = schema; if (schema_view.type == NANOARROW_TYPE_DICTIONARY) { const struct ArrowIpcDictionaryEncoding* encoding = @@ -622,24 +633,16 @@ static ArrowErrorCode ArrowIpcEncodeField( // Add the dictionary encoding to the field FLATCC_RETURN_UNLESS_0(Field_dictionary_add(builder, dict_encoding_ref), error); - // Support dictionary values with children by encoding children from - // schema->dictionary (and add a roundtrip test for a nested value type). - // Using schema below would encode the index type's children instead and - // produce a Field whose type and children do not agree. - if (schema->dictionary->n_children != 0) { - ArrowErrorSet(error, "IPC encoding of dictionary values with children unsupported"); - return ENOTSUP; - } - - NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, schema->dictionary, error)); + value_schema = schema->dictionary; + NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, value_schema, error)); } NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeFieldType(builder, &schema_view, error)); - if (schema->n_children != 0) { + if (value_schema->n_children != 0) { FLATCC_RETURN_UNLESS_0(Field_children_start(builder), error); NANOARROW_RETURN_NOT_OK( - ArrowIpcEncodeFields(builder, schema, &ns(Field_children_push_start), + ArrowIpcEncodeFields(builder, value_schema, &ns(Field_children_push_start), &ns(Field_children_push_end), dictionary_encodings, error)); FLATCC_RETURN_UNLESS_0(Field_children_end(builder), error); } @@ -657,7 +660,7 @@ static ArrowErrorCode ArrowIpcEncodeField( static ArrowErrorCode ArrowIpcEncodeSchema( flatcc_builder_t* builder, const struct ArrowSchema* schema, const struct ArrowIpcDictionaryEncodings* dictionary_encodings, int compressed_body, - struct ArrowError* error) { + int dictionary_replacement, struct ArrowError* error) { NANOARROW_DCHECK(schema->release != NULL); if (strcmp(schema->format, "+s") != 0) { @@ -694,6 +697,10 @@ static ArrowErrorCode ArrowIpcEncodeSchema( ns(Feature_enum_t) feature = ns(Feature_COMPRESSED_BODY); FLATCC_RETURN_IF_NULL(ns(Feature_vec_push(builder, &feature)), error); } + if (dictionary_replacement && dictionary_encodings->encodings.size_bytes > 0) { + ns(Feature_enum_t) feature = ns(Feature_DICTIONARY_REPLACEMENT); + FLATCC_RETURN_IF_NULL(ns(Feature_vec_push(builder, &feature)), error); + } FLATCC_RETURN_UNLESS_0(Schema_features_end(builder), error); return NANOARROW_OK; @@ -726,7 +733,8 @@ ArrowErrorCode ArrowIpcEncoderEncodeSchema(struct ArrowIpcEncoder* encoder, NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema( builder, schema, &private->dictionary_encodings, - ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, error)); + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, + private->dictionary_replacement, error)); FLATCC_RETURN_UNLESS_0(Message_header_Schema_end(builder), error); @@ -1147,6 +1155,9 @@ static ArrowErrorCode ArrowIpcEncoderEncodeDictionaryBatch( FLATCC_RETURN_UNLESS_0(DictionaryBatch_data_end(builder), error); FLATCC_RETURN_UNLESS_0(DictionaryBatch_isDelta_add(builder, is_delta ? 1 : 0), error); FLATCC_RETURN_UNLESS_0(Message_header_DictionaryBatch_end(builder), error); + + NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeMessageMetadata(private, error)); + FLATCC_RETURN_UNLESS_0(Message_bodyLength_add(builder, buffer_encoder->body_length), error); FLATCC_RETURN_IF_NULL(ns(Message_end_as_root(builder)), error); @@ -1203,7 +1214,7 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct ArrowIpcEncoder* encoder, builder, &footer->schema, &footer->dictionaries, private->has_compressed_body || ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, - error)); + /*dictionary_replacement=*/0, error)); FLATCC_RETURN_UNLESS_0(Footer_schema_end(builder), error); const struct ArrowIpcFileBlock* blocks = diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index d17782966..fa6b37b3a 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -406,6 +406,89 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaMessageMetadata) { (KeyValues{{"schema_key", "schema_value"}})); } +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryReplacementFeature) { + 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_STRING), + NANOARROW_OK); + + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + ArrowIpcEncoderSetDictionaryReplacement(encoder.get(), /*enabled=*/1); + + struct ArrowError error; + nanoarrow::UniqueBuffer message; + ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView view; + view.data.data = message->data; + view.size_bytes = message->size_bytes; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(decoder->feature_flags, NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderNestedDictionaryValueSchema) { + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + + struct ArrowSchema* dictionary_field = schema->children[0]; + ASSERT_EQ(ArrowSchemaInitFromType(dictionary_field, NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(dictionary_field, "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(dictionary_field), NANOARROW_OK); + + struct ArrowSchema* value_schema = dictionary_field->dictionary; + ASSERT_EQ(ArrowSchemaInitFromType(value_schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(value_schema, 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(value_schema->children[0], NANOARROW_TYPE_STRING), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(value_schema->children[0], "value"), NANOARROW_OK); + + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + struct ArrowError error; + nanoarrow::UniqueBuffer message; + ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView view; + view.data.data = message->data; + view.size_bytes = message->size_bytes; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + + nanoarrow::UniqueSchema roundtrip; + ASSERT_EQ(ArrowIpcDecoderDecodeSchema(decoder.get(), roundtrip.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_NE(roundtrip->children[0]->dictionary, nullptr); + EXPECT_STREQ(roundtrip->children[0]->dictionary->format, "+s"); + ASSERT_EQ(roundtrip->children[0]->dictionary->n_children, 1); + EXPECT_STREQ(roundtrip->children[0]->dictionary->children[0]->format, "u"); + EXPECT_STREQ(roundtrip->children[0]->dictionary->children[0]->name, "value"); +} + TEST(NanoarrowIpcTest, NanoarrowIpcEncoderMessageMetadataEmpty) { nanoarrow::ipc::UniqueEncoder encoder; ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); @@ -548,6 +631,8 @@ TEST(NanoarrowIpcTest, NanoarrowIpcVisitMessageMetadataError) { TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) { nanoarrow::ipc::UniqueEncoder encoder; ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); // Build a simple Utf8 values array nanoarrow::UniqueSchema values_schema; @@ -575,6 +660,13 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) { NANOARROW_OK) << error.message; + KeyValues message_key_values{{"dictionary_key", "dictionary_value"}}; + auto message_metadata = PackMetadata(message_key_values); + ASSERT_EQ( + ArrowIpcEncoderSetMessageMetadata(encoder.get(), message_metadata.get(), &error), + NANOARROW_OK) + << error.message; + // Encode a non-delta DictionaryBatch with dictionary_id=0 nanoarrow::UniqueBuffer body_buffer; EXPECT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(), /*dictionary_id=*/0, @@ -591,6 +683,21 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) { // The encapsulated message must be non-empty and 8-byte aligned EXPECT_GT(message_buffer->size_bytes, 8); EXPECT_EQ(message_buffer->size_bytes % 8, 0); + EXPECT_EQ(DecodeMessageMetadata(message_buffer.get(), decoder.get()), + message_key_values); + + // The metadata applies to exactly one message: the next DictionaryBatch has none + message_buffer->size_bytes = 0; + body_buffer->size_bytes = 0; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(), /*dictionary_id=*/0, + /*is_delta=*/0, values_view.get(), + body_buffer.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/1, + message_buffer.get()), + NANOARROW_OK); + EXPECT_EQ(DecodeMessageMetadata(message_buffer.get(), decoder.get()), KeyValues{}); } // A record batch whose columns exercise each path of the compressed body builder: diff --git a/src/nanoarrow/ipc/files_test.cc b/src/nanoarrow/ipc/files_test.cc index 73fc9724f..cc15a5251 100644 --- a/src/nanoarrow/ipc/files_test.cc +++ b/src/nanoarrow/ipc/files_test.cc @@ -214,7 +214,7 @@ class TestFile { } ArrowErrorCode WriteNanoarrowStream(const nanoarrow::UniqueSchema& schema, - const std::vector& arrays, + std::vector& arrays, enum ArrowIpcCompressionType codec, struct ArrowBuffer* buffer, struct ArrowError* error) { @@ -226,19 +226,25 @@ class TestFile { NANOARROW_RETURN_NOT_OK(ArrowIpcWriterSetCompression( writer.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, error)); - nanoarrow::UniqueArrayView array_view; + nanoarrow::UniqueSchema schema_copy; + NANOARROW_RETURN_NOT_OK(ArrowSchemaDeepCopy(schema.get(), schema_copy.get())); + nanoarrow::UniqueArrayStream array_stream; NANOARROW_RETURN_NOT_OK( - ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), error)); - - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), error)); - for (const auto& array : arrays) { - NANOARROW_RETURN_NOT_OK( - ArrowArrayViewSetArray(array_view.get(), array.get(), error)); - - NANOARROW_RETURN_NOT_OK( - ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), error)); + ArrowBasicArrayStreamInit(array_stream.get(), schema_copy.get(), arrays.size())); + + for (size_t i = 0; i < arrays.size(); i++) { + // Preserve the decoded array for the subsequent Arrow C++ comparison while + // giving the basic stream an independently releasable shared clone. + nanoarrow::UniqueArray shared; + nanoarrow::UniqueArray clone; + NANOARROW_RETURN_NOT_OK(ArrowArrayMoveShared(arrays[i].get(), shared.get())); + ArrowErrorCode result = ArrowArrayCloneShared(shared.get(), clone.get()); + ArrowArrayMove(shared.get(), arrays[i].get()); + NANOARROW_RETURN_NOT_OK(result); + ArrowBasicArrayStreamSetArray(array_stream.get(), i, clone.get()); } - return ArrowIpcWriterWriteArrayView(writer.get(), nullptr, error); + + return ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), error); } void TestEqualsArrowCpp(const std::string& dir_prefix, @@ -529,10 +535,10 @@ INSTANTIATE_TEST_SUITE_P( TestFile::OK("generated_primitive.stream"), TestFile::OK("generated_recursive_nested.stream"), TestFile::OK("generated_union.stream"), - TestFile::ReadOnly("generated_dictionary_unsigned.stream"), - TestFile::ReadOnly("generated_dictionary.stream"), - TestFile::ReadOnly("generated_nested_dictionary.stream"), - TestFile::ReadOnly("generated_extension.stream") + TestFile::OK("generated_dictionary_unsigned.stream"), + TestFile::OK("generated_dictionary.stream"), + TestFile::OK("generated_nested_dictionary.stream"), + TestFile::OK("generated_extension.stream") // Comment to keep last line from wrapping )); @@ -581,7 +587,7 @@ INSTANTIATE_TEST_SUITE_P( TestFile::OK("0.14.1/generated_primitive.stream"), TestFile::OK("0.14.1/generated_primitive_no_batches.stream"), TestFile::OK("0.14.1/generated_primitive_zerolength.stream"), - TestFile::ReadOnly("4.0.0-shareddict/generated_shared_dict.stream"), + TestFile::OK("4.0.0-shareddict/generated_shared_dict.stream"), // cpp-21.0.0 regenerated gold files TestFile::OK("cpp-21.0.0/generated_binary.stream"), TestFile::OK("cpp-21.0.0/generated_binary_no_batches.stream"), @@ -608,10 +614,10 @@ INSTANTIATE_TEST_SUITE_P( TestFile::OK("cpp-21.0.0/generated_primitive_zerolength.stream"), TestFile::OK("cpp-21.0.0/generated_recursive_nested.stream"), TestFile::OK("cpp-21.0.0/generated_union.stream"), - TestFile::ReadOnly("cpp-21.0.0/generated_dictionary.stream"), - TestFile::ReadOnly("cpp-21.0.0/generated_dictionary_unsigned.stream"), - TestFile::ReadOnly("cpp-21.0.0/generated_extension.stream"), - TestFile::ReadOnly("cpp-21.0.0/generated_nested_dictionary.stream"), + TestFile::OK("cpp-21.0.0/generated_dictionary.stream"), + TestFile::OK("cpp-21.0.0/generated_dictionary_unsigned.stream"), + TestFile::OK("cpp-21.0.0/generated_extension.stream"), + TestFile::OK("cpp-21.0.0/generated_nested_dictionary.stream"), TestFile::NotSupported("cpp-21.0.0/generated_list_view.stream"), TestFile::NotSupported("cpp-21.0.0/generated_binary_view.stream"), TestFile::NotSupported("cpp-21.0.0/generated_run_end_encoded.stream") diff --git a/src/nanoarrow/ipc/reader.c b/src/nanoarrow/ipc/reader.c index 3348a9b60..1c66cc316 100644 --- a/src/nanoarrow/ipc/reader.c +++ b/src/nanoarrow/ipc/reader.c @@ -390,14 +390,6 @@ static int ArrowIpcArrayStreamReaderReadSchemaIfNeeded( return EINVAL; } - // ...or if it uses features we don't support - if (private_data->decoder.feature_flags & - NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT) { - ArrowErrorSet(&private_data->error, - "This stream uses unsupported feature DICTIONARY_REPLACEMENT"); - return EINVAL; - } - // Notify the decoder of buffer endianness NANOARROW_RETURN_NOT_OK_WITH_ERROR( ArrowIpcDecoderSetEndianness(&private_data->decoder, diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 9879f1b63..b6e9d7921 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -181,13 +182,58 @@ struct ArrowIpcWriterPrivate { int writing_file; int64_t bytes_written; struct ArrowIpcFooter footer; + // IDs of the dictionaries written to a file, each of which may only be written once + struct ArrowBuffer file_dictionary_ids; // Message custom_metadata held for the next RecordBatch message (not for any // DictionaryBatch messages written ahead of it) struct ArrowBuffer pending_message_metadata; int has_pending_message_metadata; + // 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); @@ -207,8 +253,10 @@ ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, private->writing_file = 0; private->bytes_written = 0; ArrowIpcFooterInit(&private->footer); + 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; @@ -227,7 +275,10 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { ArrowBufferReset(&private->body_buffer); ArrowIpcFooterReset(&private->footer); + ArrowBufferReset(&private->file_dictionary_ids); ArrowBufferReset(&private->pending_message_metadata); + ArrowIpcWriterResetViewDictionaries(private); + ArrowBufferReset(&private->view_dictionaries); ArrowFree(private); } @@ -291,6 +342,11 @@ ArrowErrorCode ArrowIpcWriterWriteSchema(struct ArrowIpcWriter* writer, struct ArrowIpcWriterPrivate* private = (struct ArrowIpcWriterPrivate*)writer->private_data; + // 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)); @@ -302,6 +358,10 @@ ArrowErrorCode ArrowIpcWriterWriteSchema(struct ArrowIpcWriter* writer, if (private->writing_file) { NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowSchemaDeepCopy(in, &private->footer.schema), error); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowIpcDictionaryEncodingsAppendSchema(&private->footer.dictionaries, + &private->footer.schema), + error); } private->bytes_written += private->buffer.size_bytes; @@ -360,22 +420,29 @@ 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; - // This check is intentionally minimal: we're allowed to write one dictionary - // batch per ID in a file but we would need to add bookkeeping to keep track - // of written IDs (and usefully a fingerprint or reference to the dictionary - // so we can check if we need to emit it again). - if (private->writing_file && - (is_delta || private->footer.dictionary_blocks.size_bytes != 0)) { - ArrowErrorSet(error, - "IPC file writing supports exactly one non-delta dictionary batch"); - return ENOTSUP; + if (private->writing_file) { + if (is_delta) { + ArrowErrorSet(error, "IPC file writing does not support delta dictionary batches"); + return ENOTSUP; + } + + const int64_t* written_ids = (const int64_t*)private->file_dictionary_ids.data; + int64_t n_written_ids = + private->file_dictionary_ids.size_bytes / (int64_t)sizeof(int64_t); + for (int64_t i = 0; i < n_written_ids; i++) { + if (written_ids[i] == dictionary_id) { + ArrowErrorSet( + error, "Arrow IPC files do not support replacement of dictionary ID %" PRId64, + dictionary_id); + return ENOTSUP; + } + } } NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0)); @@ -399,6 +466,8 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( NANOARROW_RETURN_NOT_OK_WITH_ERROR( ArrowBufferAppend(&private->footer.dictionary_blocks, &block, sizeof(block)), error); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppendInt64(&private->file_dictionary_ids, dictionary_id), error); } private->bytes_written += private->buffer.size_bytes; private->bytes_written += private->body_buffer.size_bytes; @@ -410,31 +479,433 @@ 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: +// once an array is released, the addresses of its buffers may be reused for +// different values. array_view is a view of lhs, used for its storage types. +static int ArrowIpcWriterArraysShareBuffers(const struct ArrowArrayView* array_view, + const struct ArrowArray* lhs, + const struct ArrowArray* rhs) { + if (lhs->length != rhs->length || lhs->offset != rhs->offset || + lhs->null_count != rhs->null_count || lhs->n_buffers != rhs->n_buffers || + lhs->n_children != rhs->n_children || + (lhs->dictionary == NULL) != (rhs->dictionary == NULL)) { + return 0; + } + + // The last buffer of a BinaryView or StringView array holds the sizes of its variadic + // buffers, which producers typically generate on each export, so its address says + // nothing about the values. The buffers that hold the values are compared. + int64_t n_buffers = lhs->n_buffers; + if (array_view->storage_type == NANOARROW_TYPE_BINARY_VIEW || + array_view->storage_type == NANOARROW_TYPE_STRING_VIEW) { + n_buffers--; + } + + for (int64_t i = 0; i < n_buffers; i++) { + if (lhs->buffers[i] != rhs->buffers[i]) { + return 0; + } + } + + for (int64_t i = 0; i < lhs->n_children; i++) { + if (!ArrowIpcWriterArraysShareBuffers(array_view->children[i], lhs->children[i], + rhs->children[i])) { + return 0; + } + } + + return lhs->dictionary == NULL || + ArrowIpcWriterArraysShareBuffers(array_view->dictionary, lhs->dictionary, + rhs->dictionary); +} + // Walk the array in the same depth-first order the schema encoder uses to assign // dictionary ids (see ArrowIpcDictionaryEncodingsAppendSchema): a dictionary-encoded // node claims the next id before descending into its children and then its values. -// Emitting a full (non-delta) DictionaryBatch for each dictionary before every -// RecordBatch keeps each batch's indices valid against the dictionary that precedes -// it, which is required because each array in the stream carries its own dictionary. -// In the future we can reduce the number of dictionaries emitted by checking for -// identical dictionary arrays. -static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( +// A dictionary is written after any dictionaries nested in its values so that they +// are available when a reader decodes it. A dictionary that shares its buffers with +// the corresponding dictionary of the previous array (if any) was already written; +// otherwise, it is written as a full (non-delta) DictionaryBatch that replaces the +// previous one. array and previous may both be NULL, in which case every dictionary +// is written. +static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArray( + struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view, + const struct ArrowArray* array, const struct ArrowArray* previous, int64_t* next_id, + 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(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view->children[i], array == NULL ? NULL : array->children[i], + previous == NULL ? NULL : previous->children[i], next_id, error)); + } + + if (array_view->dictionary != NULL) { + const struct ArrowArray* previous_dictionary = + previous == NULL ? NULL : previous->dictionary; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view->dictionary, array == NULL ? NULL : array->dictionary, + previous_dictionary, next_id, error)); + + if (previous_dictionary == NULL || + !ArrowIpcWriterArraysShareBuffers(array_view->dictionary, array->dictionary, + previous_dictionary)) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( + writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error)); + } + } + + 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 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, 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, length); + } + if (result == NANOARROW_OK) { + result = ArrowArrayAppendStorageFromArrayView(out, &slice, 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; +} + +// 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). 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, + 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; + } + } + + // 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, 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); + *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, struct ArrowError* error) { + int64_t* next_id, int* written, struct ArrowError* error) { + int64_t dictionary_id = *next_id; if (array_view->dictionary != NULL) { - int64_t dictionary_id = (*next_id)++; - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( - writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error)); + (*next_id)++; } for (int64_t i = 0; i < array_view->n_children; i++) { - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( - writer, array_view->children[i], next_id, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, array_view->children[i], next_id, written, error)); } if (array_view->dictionary != NULL) { - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( - writer, array_view->dictionary, next_id, error)); + 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; @@ -445,14 +916,15 @@ ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( 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)); + int written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, in, &next_dictionary_id, &written, error)); return ArrowIpcWriterWriteArrayView(writer, in, error); } static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl( struct ArrowIpcWriter* writer, struct ArrowArrayStream* in, - struct ArrowSchema* schema, struct ArrowArray* array, + struct ArrowSchema* schema, struct ArrowArray* array, struct ArrowArray* previous, struct ArrowArrayView* array_view, struct ArrowError* error) { NANOARROW_RETURN_NOT_OK(ArrowArrayStreamGetSchema(in, schema, error)); NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteSchema(writer, schema, error)); @@ -466,9 +938,25 @@ static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl( NANOARROW_RETURN_NOT_OK(ArrowArrayViewSetArray(array_view, array, error)); - NANOARROW_RETURN_NOT_OK( - ArrowIpcWriterWriteArrayViewWithDictionaries(writer, array_view, error)); - ArrowArrayRelease(array); + int64_t next_dictionary_id = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view, array, previous->release == NULL ? NULL : previous, + &next_dictionary_id, error)); + + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteArrayView(writer, array_view, error)); + + // Keep an array with dictionaries alive until the next array is written so that + // its dictionaries' buffers can't be released and their addresses reused for + // different values while we compare them by address. + if (previous->release != NULL) { + ArrowArrayRelease(previous); + } + + if (next_dictionary_id > 0) { + ArrowArrayMove(array, previous); + } else { + ArrowArrayRelease(array); + } } // The stream is complete, signal the end to the caller @@ -482,11 +970,12 @@ ArrowErrorCode ArrowIpcWriterWriteArrayStream(struct ArrowIpcWriter* writer, struct ArrowSchema schema = {.release = NULL}; struct ArrowArray array = {.release = NULL}; + struct ArrowArray previous = {.release = NULL}; struct ArrowArrayView array_view; ArrowArrayViewInitFromType(&array_view, NANOARROW_TYPE_UNINITIALIZED); - ArrowErrorCode result = - ArrowIpcWriterWriteArrayStreamImpl(writer, in, &schema, &array, &array_view, error); + ArrowErrorCode result = ArrowIpcWriterWriteArrayStreamImpl( + writer, in, &schema, &array, &previous, &array_view, error); if (schema.release != NULL) { ArrowSchemaRelease(&schema); @@ -496,6 +985,10 @@ ArrowErrorCode ArrowIpcWriterWriteArrayStream(struct ArrowIpcWriter* writer, ArrowArrayRelease(&array); } + if (previous.release != NULL) { + ArrowArrayRelease(&previous); + } + ArrowArrayViewReset(&array_view); return result; diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index c01e09986..35fa9487e 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -18,11 +18,18 @@ #include #include +#include #include #include #include +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) +#include +#include +#include +#endif + #include "nanoarrow/nanoarrow_ipc.hpp" TEST(NanoarrowIpcWriter, OutputStreamBuffer) { @@ -281,28 +288,32 @@ TEST(NanoarrowIpcWriter, WriteDictionaryBatch) { // one block tracked in file mode EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); + // a dictionary ID can't be written to a file twice int64_t bytes_written = p2->bytes_written; EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(), /*dictionary_id=*/0, /*is_delta=*/0, values_view.get(), &error), ENOTSUP); EXPECT_STREQ(error.message, - "IPC file writing supports exactly one non-delta dictionary batch"); + "Arrow IPC files do not support replacement of dictionary ID 0"); EXPECT_EQ(p2->bytes_written, bytes_written); EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); - nanoarrow::ipc::UniqueOutputStream stream3; - nanoarrow::UniqueBuffer output3; - ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream3.get(), output3.get()), NANOARROW_OK); - - nanoarrow::ipc::UniqueWriter writer3; - ASSERT_EQ(ArrowIpcWriterInit(writer3.get(), stream3.get()), NANOARROW_OK); - ASSERT_EQ(ArrowIpcWriterStartFile(writer3.get(), &error), NANOARROW_OK) + // ...but other dictionary IDs can + EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(), /*dictionary_id=*/1, + /*is_delta=*/0, values_view.get(), &error), + NANOARROW_OK) << error.message; - EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer3.get(), /*dictionary_id=*/0, + EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, + 2 * sizeof(struct ArrowIpcFileBlock)); + + // delta dictionaries aren't supported in files + bytes_written = p2->bytes_written; + EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(), /*dictionary_id=*/2, /*is_delta=*/1, values_view.get(), &error), ENOTSUP); EXPECT_STREQ(error.message, - "IPC file writing supports exactly one non-delta dictionary batch"); + "IPC file writing does not support delta dictionary batches"); + EXPECT_EQ(p2->bytes_written, bytes_written); } // Build a struct array with a single dictionary-encoded (int32 -> utf8) child. @@ -335,6 +346,523 @@ static void MakeDictionaryStructArray(struct ArrowArray* array, ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); } +static std::vector DecodeMessageTypes(const struct ArrowBuffer* buffer) { + std::vector message_types; + struct ArrowBufferView remaining; + remaining.data.as_uint8 = buffer->data; + remaining.size_bytes = buffer->size_bytes; + struct ArrowIpcDecoder decoder; + struct ArrowError error; + ArrowIpcDecoderInit(&decoder); + + while (remaining.size_bytes > 0) { + int result = ArrowIpcDecoderVerifyHeader(&decoder, remaining, &error); + if (result == ENODATA) { + break; + } + + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + + message_types.push_back(decoder.message_type); + int64_t message_size = ((decoder.header_size_bytes + 7) / 8) * 8 + + ((decoder.body_size_bytes + 7) / 8) * 8; + EXPECT_LE(message_size, remaining.size_bytes); + if (message_size > remaining.size_bytes) { + break; + } + + remaining.data.as_uint8 += message_size; + remaining.size_bytes -= message_size; + } + + ArrowIpcDecoderReset(&decoder); + return message_types; +} + +static std::vector DecodeDictionaryIds(const struct ArrowBuffer* buffer) { + std::vector dictionary_ids; + struct ArrowBufferView remaining; + remaining.data.as_uint8 = buffer->data; + remaining.size_bytes = buffer->size_bytes; + struct ArrowIpcDecoder decoder; + struct ArrowError error; + ArrowIpcDecoderInit(&decoder); + + while (remaining.size_bytes > 0) { + int result = ArrowIpcDecoderVerifyHeader(&decoder, 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, remaining, &error); + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + dictionary_ids.push_back(decoder.dictionary->id); + } + + int64_t message_size = ((decoder.header_size_bytes + 7) / 8) * 8 + + ((decoder.body_size_bytes + 7) / 8) * 8; + remaining.data.as_uint8 += message_size; + remaining.size_bytes -= message_size; + } + + ArrowIpcDecoderReset(&decoder); + return dictionary_ids; +} + +static void MakeNestedDictionaryStructArray(struct ArrowArray* array, + struct ArrowSchema* schema, + const char* inner_value1 = "bar") { + ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); + + struct ArrowSchema* outer_field = schema->children[0]; + ASSERT_EQ(ArrowSchemaInitFromType(outer_field, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(outer_field, "outer"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(outer_field), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(outer_field->dictionary, NANOARROW_TYPE_STRUCT), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(outer_field->dictionary, 1), NANOARROW_OK); + + struct ArrowSchema* inner_field = outer_field->dictionary->children[0]; + ASSERT_EQ(ArrowSchemaInitFromType(inner_field, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(inner_field, "inner"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(inner_field), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(inner_field->dictionary, NANOARROW_TYPE_STRING), + NANOARROW_OK); + + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK); + struct ArrowArray* outer_indices = array->children[0]; + struct ArrowArray* outer_values = outer_indices->dictionary; + struct ArrowArray* inner_indices = outer_values->children[0]; + struct ArrowArray* inner_values = inner_indices->dictionary; + + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(inner_values, ArrowCharView("foo")), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(inner_values, ArrowCharView(inner_value1)), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(inner_indices, 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(inner_indices, 1), NANOARROW_OK); + outer_values->length = 2; + ASSERT_EQ(ArrowArrayAppendInt(outer_indices, 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(outer_indices, 1), NANOARROW_OK); + array->length = 2; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); +} + +TEST(NanoarrowIpcWriter, WritesNestedDictionariesDependencyFirst) { + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + MakeNestedDictionaryStructArray(array.get(), schema.get()); + + nanoarrow::UniqueArrayStream array_stream; + ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), schema.get(), 1), NANOARROW_OK); + ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get()); + + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream out_stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error), + NANOARROW_OK) + << error.message; + + EXPECT_EQ(DecodeDictionaryIds(output.get()), (std::vector{1, 0})); + + 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; + nanoarrow::UniqueArray roundtrip_array; + ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(), &error), + NANOARROW_OK) + << error.message; + + struct ArrowArray* outer_dictionary = roundtrip_array->children[0]->dictionary; + ASSERT_NE(outer_dictionary, nullptr); + ASSERT_EQ(outer_dictionary->n_children, 1); + struct ArrowArray* inner_dictionary = outer_dictionary->children[0]->dictionary; + ASSERT_NE(inner_dictionary, nullptr); + EXPECT_EQ(inner_dictionary->length, 2); + + nanoarrow::UniqueArrayView inner_view; + ArrowArrayViewInitFromType(inner_view.get(), NANOARROW_TYPE_STRING); + ASSERT_EQ(ArrowArrayViewSetArray(inner_view.get(), inner_dictionary, &error), + NANOARROW_OK); + EXPECT_EQ(ArrowArrayViewGetStringUnsafe(inner_view.get(), 0), ArrowCharView("foo")); + EXPECT_EQ(ArrowArrayViewGetStringUnsafe(inner_view.get(), 1), ArrowCharView("bar")); +} + +// Write arrays using ArrowIpcWriterWriteArrayStream() (consuming schema and arrays) +static ArrowErrorCode WriteArrays(struct ArrowSchema* schema, + std::vector* arrays, + struct ArrowBuffer* output, bool as_file, + struct ArrowError* error) { + nanoarrow::UniqueArrayStream array_stream; + NANOARROW_RETURN_NOT_OK( + ArrowBasicArrayStreamInit(array_stream.get(), schema, arrays->size())); + for (size_t i = 0; i < arrays->size(); i++) { + ArrowBasicArrayStreamSetArray(array_stream.get(), i, (*arrays)[i].get()); + } + + nanoarrow::ipc::UniqueOutputStream out_stream; + NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output)); + nanoarrow::ipc::UniqueWriter writer; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterInit(writer.get(), out_stream.get())); + if (as_file) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterStartFile(writer.get(), error)); + } + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), error)); + if (as_file) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterFinalizeFile(writer.get(), error)); + } + return NANOARROW_OK; +} + +// Read all arrays of a stream (taking ownership of buffer) +static void ReadArrays(struct ArrowBuffer* buffer, + std::vector* arrays) { + struct ArrowError error; + struct ArrowIpcInputStream input; + ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, buffer), NANOARROW_OK); + nanoarrow::UniqueArrayStream reader; + ASSERT_EQ(ArrowIpcArrayStreamReaderInit(reader.get(), &input, nullptr), NANOARROW_OK); + while (true) { + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), array.get(), &error), NANOARROW_OK) + << error.message; + if (array->release == nullptr) { + break; + } + arrays->push_back(std::move(array)); + } +} + +static std::vector StringValues(const struct ArrowArray* array) { + std::vector values; + nanoarrow::UniqueArrayView view; + ArrowArrayViewInitFromType(view.get(), NANOARROW_TYPE_STRING); + EXPECT_EQ(ArrowArrayViewSetArray(view.get(), array, nullptr), NANOARROW_OK); + for (int64_t i = 0; i < view->length; i++) { + struct ArrowStringView value = ArrowArrayViewGetStringUnsafe(view.get(), i); + values.emplace_back(value.data, value.size_bytes); + } + return values; +} + +// Create count arrays that share all buffers with array (consuming it) +static std::vector CloneShared(struct ArrowArray* array, + size_t count) { + nanoarrow::UniqueArray shared; + EXPECT_EQ(ArrowArrayMoveShared(array, shared.get()), NANOARROW_OK); + std::vector arrays(count); + for (auto& clone : arrays) { + EXPECT_EQ(ArrowArrayCloneShared(shared.get(), clone.get()), NANOARROW_OK); + } + return arrays; +} + +TEST(NanoarrowIpcWriter, DoesNotRepeatSharedDictionary) { + struct ArrowError error; + + for (bool as_file : {false, true}) { + SCOPED_TRACE(as_file ? "file" : "stream"); + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + MakeDictionaryStructArray(array.get(), schema.get()); + std::vector arrays = CloneShared(array.get(), 3); + + nanoarrow::UniqueBuffer output; + ASSERT_EQ(WriteArrays(schema.get(), &arrays, output.get(), as_file, &error), + NANOARROW_OK) + << error.message; + + if (as_file) { +#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(); + auto arrow_reader = maybe_arrow_reader.ValueUnsafe(); + ASSERT_EQ(arrow_reader->num_record_batches(), 3); + auto maybe_batch = arrow_reader->ReadRecordBatch(2); + ASSERT_TRUE(maybe_batch.ok()) << maybe_batch.status(); + auto arrow_dictionary = std::static_pointer_cast( + maybe_batch.ValueUnsafe()->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + ASSERT_EQ(arrow_values->length(), 2); + EXPECT_EQ(arrow_values->GetString(1), "bar"); +#endif + continue; + } + + 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_RECORD_BATCH})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 3); + for (const auto& roundtrip_array : roundtrip) { + EXPECT_EQ(StringValues(roundtrip_array->children[0]->dictionary), + (std::vector{"foo", "bar"})); + } + } +} + +TEST(NanoarrowIpcWriter, EmitsChangedDictionary) { + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueSchema unused_schema; + std::vector arrays(2); + MakeDictionaryStructArray(arrays[0].get(), schema.get()); + MakeDictionaryStructArray(arrays[1].get(), unused_schema.get(), "baz"); + + nanoarrow::UniqueBuffer output; + ASSERT_EQ(WriteArrays(schema.get(), &arrays, output.get(), /*as_file=*/false, &error), + NANOARROW_OK) + << error.message; + 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})); + +#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(); + auto arrow_reader = maybe_arrow_reader.ValueUnsafe(); + std::shared_ptr arrow_batch; + ASSERT_TRUE(arrow_reader->ReadNext(&arrow_batch).ok()); + ASSERT_TRUE(arrow_reader->ReadNext(&arrow_batch).ok()); + 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), "baz"); +#endif + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); + EXPECT_EQ(StringValues(roundtrip[0]->children[0]->dictionary), + (std::vector{"foo", "bar"})); + EXPECT_EQ(StringValues(roundtrip[1]->children[0]->dictionary), + (std::vector{"foo", "baz"})); +} + +TEST(NanoarrowIpcWriter, RejectsChangedDictionaryInFile) { + struct ArrowError error; + nanoarrow::UniqueSchema schema; + std::vector arrays(2); + MakeDictionaryStructArray(arrays[0].get(), schema.get()); + nanoarrow::UniqueSchema unused_schema; + MakeDictionaryStructArray(arrays[1].get(), unused_schema.get(), "baz"); + + nanoarrow::UniqueBuffer output; + EXPECT_EQ(WriteArrays(schema.get(), &arrays, output.get(), /*as_file=*/true, &error), + ENOTSUP); + EXPECT_STREQ(error.message, + "Arrow IPC files do not support replacement of dictionary ID 0"); +} + +TEST(NanoarrowIpcWriter, DoesNotRepeatSharedNestedDictionaries) { + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + MakeNestedDictionaryStructArray(array.get(), schema.get()); + std::vector arrays = CloneShared(array.get(), 3); + + // Replace the inner dictionary of the last array: both the inner and the outer + // dictionary (whose values reference the inner one) must be written again + struct ArrowArray* inner_dictionary = + arrays[2]->children[0]->dictionary->children[0]->dictionary; + ArrowArrayRelease(inner_dictionary); + ASSERT_EQ(ArrowArrayInitFromType(inner_dictionary, NANOARROW_TYPE_STRING), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(inner_dictionary), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(inner_dictionary, ArrowCharView("foo")), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(inner_dictionary, ArrowCharView("baz")), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(inner_dictionary, &error), NANOARROW_OK); + + nanoarrow::UniqueBuffer output; + ASSERT_EQ(WriteArrays(schema.get(), &arrays, output.get(), /*as_file=*/false, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(DecodeDictionaryIds(output.get()), (std::vector{1, 0, 1, 0})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 3); + EXPECT_EQ(StringValues(roundtrip[1]->children[0]->dictionary->children[0]->dictionary), + (std::vector{"foo", "bar"})); + EXPECT_EQ(StringValues(roundtrip[2]->children[0]->dictionary->children[0]->dictionary), + (std::vector{"foo", "baz"})); +} + +// An ArrowArrayStream whose producer writes each array into the memory of an array +// that was already released (if any), like a producer that recycles buffers. Each +// array has a different dictionary, but a released dictionary's buffer addresses are +// reused for the next one. +struct RecyclingSlot { + bool in_use; + int32_t indices[2]; + int32_t offsets[3]; + char data[4]; + const void* struct_buffers[1]; + const void* indices_buffers[2]; + const void* dictionary_buffers[3]; + struct ArrowArray dictionary; + struct ArrowArray indices_array; + struct ArrowArray* children[1]; +}; + +struct RecyclingStream { + RecyclingSlot slots[2]{}; + std::vector dictionaries{"aabb", "ccdd", "eeff"}; + std::vector slots_used; +}; + +static void RecyclingReleaseChild(struct ArrowArray* array) { array->release = nullptr; } + +static void RecyclingRelease(struct ArrowArray* array) { + static_cast(array->private_data)->in_use = false; + array->release = nullptr; +} + +static int RecyclingGetSchema(struct ArrowArrayStream*, struct ArrowSchema* out) { + NANOARROW_RETURN_NOT_OK(ArrowSchemaInitFromType(out, NANOARROW_TYPE_STRUCT)); + NANOARROW_RETURN_NOT_OK(ArrowSchemaAllocateChildren(out, 1)); + NANOARROW_RETURN_NOT_OK( + ArrowSchemaInitFromType(out->children[0], NANOARROW_TYPE_INT32)); + NANOARROW_RETURN_NOT_OK(ArrowSchemaSetName(out->children[0], "dict_col")); + NANOARROW_RETURN_NOT_OK(ArrowSchemaAllocateDictionary(out->children[0])); + return ArrowSchemaInitFromType(out->children[0]->dictionary, NANOARROW_TYPE_STRING); +} + +static int RecyclingGetNext(struct ArrowArrayStream* stream, struct ArrowArray* out) { + auto* private_data = static_cast(stream->private_data); + size_t i = private_data->slots_used.size(); + if (i == private_data->dictionaries.size()) { + out->release = nullptr; + return NANOARROW_OK; + } + + int slot_index = private_data->slots[0].in_use ? 1 : 0; + RecyclingSlot* slot = &private_data->slots[slot_index]; + if (slot->in_use) { + return EBUSY; + } + private_data->slots_used.push_back(slot_index); + + memcpy(slot->data, private_data->dictionaries[i].data(), sizeof(slot->data)); + slot->offsets[0] = 0; + slot->offsets[1] = 2; + slot->offsets[2] = 4; + slot->indices[0] = 0; + slot->indices[1] = 1; + + slot->dictionary_buffers[0] = nullptr; + slot->dictionary_buffers[1] = slot->offsets; + slot->dictionary_buffers[2] = slot->data; + slot->dictionary = {}; + slot->dictionary.length = 2; + slot->dictionary.n_buffers = 3; + slot->dictionary.buffers = slot->dictionary_buffers; + slot->dictionary.release = &RecyclingReleaseChild; + + slot->indices_buffers[0] = nullptr; + slot->indices_buffers[1] = slot->indices; + slot->indices_array = {}; + slot->indices_array.length = 2; + slot->indices_array.n_buffers = 2; + slot->indices_array.buffers = slot->indices_buffers; + slot->indices_array.dictionary = &slot->dictionary; + slot->indices_array.release = &RecyclingReleaseChild; + + slot->struct_buffers[0] = nullptr; + slot->children[0] = &slot->indices_array; + *out = {}; + out->length = 2; + out->n_buffers = 1; + out->buffers = slot->struct_buffers; + out->n_children = 1; + out->children = slot->children; + out->release = &RecyclingRelease; + out->private_data = slot; + slot->in_use = true; + return NANOARROW_OK; +} + +static const char* RecyclingGetLastError(struct ArrowArrayStream*) { return nullptr; } + +static void RecyclingStreamRelease(struct ArrowArrayStream* stream) { + stream->release = nullptr; +} + +TEST(NanoarrowIpcWriter, DoesNotRepeatDictionaryAtRecycledAddress) { + struct ArrowError error; + RecyclingStream private_data; + struct ArrowArrayStream array_stream; + array_stream.get_schema = &RecyclingGetSchema; + array_stream.get_next = &RecyclingGetNext; + array_stream.get_last_error = &RecyclingGetLastError; + array_stream.release = &RecyclingStreamRelease; + array_stream.private_data = &private_data; + + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream out_stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), &array_stream, &error), + NANOARROW_OK) + << error.message; + ArrowArrayStreamRelease(&array_stream); + + // The writer keeps the previous array alive while requesting the next one, so a + // dictionary that it compares by address is never at a recycled address + EXPECT_EQ(private_data.slots_used, (std::vector{0, 1, 0})); + EXPECT_FALSE(private_data.slots[0].in_use); + EXPECT_FALSE(private_data.slots[1].in_use); + EXPECT_EQ(DecodeDictionaryIds(output.get()), (std::vector{0, 0, 0})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 3); + EXPECT_EQ(StringValues(roundtrip[0]->children[0]->dictionary), + (std::vector{"aa", "bb"})); + EXPECT_EQ(StringValues(roundtrip[1]->children[0]->dictionary), + (std::vector{"cc", "dd"})); + EXPECT_EQ(StringValues(roundtrip[2]->children[0]->dictionary), + (std::vector{"ee", "ff"})); +} + // Write a dictionary-encoded stream through the high-level WriteArrayStream path // and read it back through the IPC reader, confirming the DictionaryBatch is // emitted automatically and the decoded values match. @@ -360,6 +888,16 @@ TEST(NanoarrowIpcWriter, RoundtripDictionaryStream) { NANOARROW_OK) << error.message; + // The schema declares that dictionaries may be replaced + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView data; + data.data.data = output->data; + data.size_bytes = output->size_bytes; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), data, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(decoder->feature_flags, NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT); + // Read the encoded bytes back struct ArrowIpcInputStream input; ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, output.get()), NANOARROW_OK); @@ -414,6 +952,55 @@ TEST(NanoarrowIpcWriter, RoundtripDictionaryStream) { EXPECT_EQ(roundtrip_array->release, nullptr); } +// Write a dictionary-encoded file through the WriteArrayStream path, which +// requires the footer schema to resolve its dictionary IDs. +TEST(NanoarrowIpcWriter, RoundtripDictionaryFile) { + struct ArrowError error; + + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + MakeDictionaryStructArray(array.get(), schema.get()); + + nanoarrow::UniqueArrayStream array_stream; + ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), schema.get(), 1), NANOARROW_OK); + ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get()); + + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream out_stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()), NANOARROW_OK); + + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterStartFile(writer.get(), &error), NANOARROW_OK) << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK) + << error.message; + + auto* p = static_cast(writer->private_data); + EXPECT_EQ(p->footer.dictionary_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); + EXPECT_EQ(p->footer.record_batch_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); + +#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(); + auto arrow_reader = maybe_arrow_reader.ValueUnsafe(); + ASSERT_EQ(arrow_reader->num_record_batches(), 1); + auto maybe_batch = arrow_reader->ReadRecordBatch(0); + ASSERT_TRUE(maybe_batch.ok()) << maybe_batch.status(); + auto arrow_dictionary = std::static_pointer_cast( + maybe_batch.ValueUnsafe()->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + ASSERT_EQ(arrow_values->length(), 2); + EXPECT_EQ(arrow_values->GetString(0), "foo"); + EXPECT_EQ(arrow_values->GetString(1), "bar"); +#endif +} + // A struct array with a single int32 column of repeating values (i.e., compressible) static constexpr int64_t kCompressibleBatchLength = 1024; @@ -892,6 +1479,549 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) { } } +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesInFile) { + // 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; + 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 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(), schema.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &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); +} + +// 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 a607df80b..8c40daf62 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -108,6 +108,8 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetCompression) #define ArrowIpcEncoderSetCompressor \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetCompressor) +#define ArrowIpcEncoderSetDictionaryReplacement \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetDictionaryReplacement) #define ArrowIpcEncoderEncodeSchema \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderEncodeSchema) #define ArrowIpcEncoderEncodeSimpleRecordBatch \ @@ -991,9 +993,10 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderFinalizeBuffer( /// \brief Set the custom metadata of the next encoded message /// -/// Attaches metadata to the next message encoded by ArrowIpcEncoderEncodeSchema() or -/// ArrowIpcEncoderEncodeSimpleRecordBatch() (i.e., Message::custom_metadata, which is -/// distinct from the metadata of the Schema or Field that the message may contain). +/// Attaches metadata to the next message encoded by ArrowIpcEncoderEncodeSchema(), +/// ArrowIpcEncoderEncodeSimpleRecordBatch(), or +/// ArrowIpcEncoderEncodeSimpleDictionaryBatch() (i.e., Message::custom_metadata, which +/// is distinct from the metadata of the Schema or Field that the message may contain). /// The metadata applies to exactly one message: after a message is encoded the /// encoder's message metadata is cleared. Any metadata that was set but not yet /// encoded is replaced by this call; pass NULL to clear it. @@ -1046,6 +1049,15 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompressor( struct ArrowIpcEncoder* encoder, struct ArrowIpcCompressor* compressor); +/// \brief Declare dictionary replacement support in subsequently encoded schemas +/// +/// Enable this before encoding the schema of a stream that may contain more than one +/// non-delta DictionaryBatch with the same dictionary ID. The +/// DICTIONARY_REPLACEMENT feature is only written if the schema contains at least one +/// dictionary-encoded field. This option is disabled by default. +NANOARROW_DLL void ArrowIpcEncoderSetDictionaryReplacement( + struct ArrowIpcEncoder* encoder, char enabled); + /// \brief Encode an ArrowSchema /// /// Returns ENOMEM if allocation fails, NANOARROW_OK otherwise. @@ -1193,13 +1205,24 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// \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. +/// by the RecordBatch itself, 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, and nested dictionaries +/// are written before the dictionaries that refer to them. +/// +/// 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. 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. /// /// Errors are propagated from the underlying encoder and output byte stream. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( @@ -1211,6 +1234,11 @@ 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: 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). /// /// Errors are propagated from the underlying encoder and output byte stream. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( @@ -1219,6 +1247,17 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( /// \brief Write an entire stream (including EOS) to the output byte stream /// +/// Dictionaries are written before the first array that references them. For +/// subsequent arrays, a dictionary is written again (as a replacement) unless it has +/// the same length, offset, and null count and points to the same buffers as the +/// corresponding dictionary of the previous array (e.g., because the producer shares +/// one dictionary among all arrays of the stream). The variadic buffer sizes of a +/// BinaryView or StringView array are not compared, since producers may generate them +/// on each export. To ensure that buffer addresses +/// are not reused for different values, the previous array is kept alive until the +/// next array is written. Buffers must not be modified while they are referenced by +/// an exported array. When writing a file, a changed dictionary returns ENOTSUP. +/// /// Errors are propagated from the underlying encoder, array stream, and output byte /// stream. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayStream(struct ArrowIpcWriter* writer,