diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index ab7f0b60d..1b0c8dcf9 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); @@ -555,6 +565,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 = @@ -623,24 +634,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); } @@ -658,7 +661,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) { @@ -695,6 +698,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; @@ -727,7 +734,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); @@ -1124,6 +1132,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); @@ -1180,7 +1191,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 00fd278e5..2020d94c4 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -326,6 +326,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); @@ -468,6 +551,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; @@ -495,6 +580,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, @@ -511,6 +603,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 c699ca331..1cef35d5f 100644 --- a/src/nanoarrow/ipc/reader.c +++ b/src/nanoarrow/ipc/reader.c @@ -389,14 +389,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 bb8c79710..e4f58a932 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,6 +182,8 @@ 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; }; ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, @@ -202,6 +205,7 @@ ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, private->writing_file = 0; private->bytes_written = 0; ArrowIpcFooterInit(&private->footer); + ArrowBufferInit(&private->file_dictionary_ids); writer->private_data = private; return NANOARROW_OK; @@ -220,6 +224,7 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { ArrowBufferReset(&private->body_buffer); ArrowIpcFooterReset(&private->footer); + ArrowBufferReset(&private->file_dictionary_ids); ArrowFree(private); } @@ -265,6 +270,9 @@ 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); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0)); NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSchema(&private->encoder, in, error)); @@ -276,6 +284,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; @@ -335,15 +347,23 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( 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 EINVAL; + } + } } NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0)); @@ -367,6 +387,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; @@ -378,31 +400,71 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( return NANOARROW_OK; } +// 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. +static int ArrowIpcWriterArraysShareBuffers(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; + } + + for (int64_t i = 0; i < lhs->n_buffers; i++) { + if (lhs->buffers[i] != rhs->buffers[i]) { + return 0; + } + } + + for (int64_t i = 0; i < lhs->n_children; i++) { + if (!ArrowIpcWriterArraysShareBuffers(lhs->children[i], rhs->children[i])) { + return 0; + } + } + + return lhs->dictionary == NULL || + ArrowIpcWriterArraysShareBuffers(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. +static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArray( struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view, - int64_t* next_id, struct ArrowError* error) { + 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) { - 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(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view->children[i], array->children[i], + previous == NULL ? NULL : previous->children[i], next_id, error)); } if (array_view->dictionary != NULL) { - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( - writer, array_view->dictionary, next_id, error)); + const struct ArrowArray* previous_dictionary = + previous == NULL ? NULL : previous->dictionary; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view->dictionary, array->dictionary, previous_dictionary, next_id, + error)); + + if (previous_dictionary == NULL || + !ArrowIpcWriterArraysShareBuffers(array->dictionary, previous_dictionary)) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( + writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error)); + } } return NANOARROW_OK; @@ -410,7 +472,7 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( 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)); @@ -425,11 +487,24 @@ static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl( NANOARROW_RETURN_NOT_OK(ArrowArrayViewSetArray(array_view, array, error)); int64_t next_dictionary_id = 0; - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( - writer, array_view, &next_dictionary_id, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( + writer, array_view, array, previous->release == NULL ? NULL : previous, + &next_dictionary_id, error)); NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteArrayView(writer, array_view, error)); - ArrowArrayRelease(array); + + // 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 @@ -443,11 +518,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); @@ -457,6 +533,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 3b08d0c89..f542c34ce 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -18,6 +18,16 @@ #include #include +#include + +#include +#include + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) +#include +#include +#include +#endif #include "nanoarrow/nanoarrow_ipc.hpp" @@ -277,33 +287,38 @@ 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); + EINVAL); 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. static void MakeDictionaryStructArray(struct ArrowArray* array, - struct ArrowSchema* schema) { + struct ArrowSchema* schema, + const char* value1 = "bar") { ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), @@ -320,7 +335,7 @@ static void MakeDictionaryStructArray(struct ArrowArray* array, ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("foo")), NANOARROW_OK); - ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("bar")), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView(value1)), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK); ASSERT_EQ(ArrowArrayAppendInt(indices, 1), NANOARROW_OK); @@ -330,6 +345,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), + EINVAL); + 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. @@ -354,6 +886,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); @@ -406,6 +948,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; diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 7f5559759..d45cc1bc8 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 \ @@ -979,9 +981,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. @@ -1034,6 +1037,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. @@ -1168,6 +1180,10 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// dictionary_id must match the id assigned to the dictionary-encoded field in the /// schema. is_delta selects DictionaryBatch.isDelta. values_view must not itself be /// dictionary-encoded. The writer does not check that a schema was already written. +/// values_view is always written: the writer only sees a borrowed view and cannot +/// check whether it is identical to a dictionary that was already written. When +/// writing a file, returns ENOTSUP for a delta and EINVAL 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( @@ -1176,6 +1192,15 @@ 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). 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 EINVAL. +/// /// Errors are propagated from the underlying encoder, array stream, and output byte /// stream. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayStream(struct ArrowIpcWriter* writer,