From a34ce843042edd9abc52794aca13a49b3fd76896 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:42:50 -0400 Subject: [PATCH 01/10] fix(ipc): attach message metadata to dictionary batches ArrowIpcEncoderSetMessageMetadata() was ignored by ArrowIpcEncoderEncodeSimpleDictionaryBatch(); encode it like the schema and record batch messages do. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/encoder.c | 3 +++ src/nanoarrow/ipc/encoder_test.cc | 24 ++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 7 ++++--- 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index ab7f0b60d..ffc9d6f8e 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -1124,6 +1124,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); diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 00fd278e5..6457e09ea 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -468,6 +468,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 +497,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 +520,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/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 7f5559759..87ee87cc6 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -979,9 +979,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. From e5922b337465bdaa980eab3109fd2db14b186d0a Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:43:06 -0400 Subject: [PATCH 02/10] fix(ipc): encode children of dictionary value fields The IPC Field for a dictionary-encoded column describes the dictionary value type, so its children must come from schema->dictionary rather than from the index schema. This replaces the ENOTSUP for dictionary values with children (e.g., struct dictionaries whose children are themselves dictionary-encoded). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/encoder.c | 17 ++++------- src/nanoarrow/ipc/encoder_test.cc | 48 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index ffc9d6f8e..c2f281f83 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -555,6 +555,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 +624,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); } diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 6457e09ea..b5eadd23e 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -326,6 +326,54 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaMessageMetadata) { (KeyValues{{"schema_key", "schema_value"}})); } +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); From ad860f94c61a69aa6dd169afcd119bfb80a7fcd7 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:43:34 -0400 Subject: [PATCH 03/10] feat(ipc): advertise DICTIONARY_REPLACEMENT in written streams Add ArrowIpcEncoderSetDictionaryReplacement() to declare the DICTIONARY_REPLACEMENT feature in an encoded Schema message. The writer enables it in stream mode, where a dictionary ID may be written more than once, and the reader no longer rejects streams that declare it (the decoder already replaces dictionaries). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/encoder.c | 21 ++++++++++++++++--- src/nanoarrow/ipc/encoder_test.cc | 35 +++++++++++++++++++++++++++++++ src/nanoarrow/ipc/reader.c | 8 ------- src/nanoarrow/ipc/writer.c | 3 +++ src/nanoarrow/ipc/writer_test.cc | 10 +++++++++ src/nanoarrow/nanoarrow_ipc.h | 11 ++++++++++ 6 files changed, 77 insertions(+), 11 deletions(-) diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index c2f281f83..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); @@ -651,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) { @@ -688,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; @@ -720,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); @@ -1176,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 b5eadd23e..2020d94c4 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -326,6 +326,41 @@ 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); 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..65c0e449e 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -265,6 +265,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)); diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 3b08d0c89..fba040fa8 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -354,6 +354,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); diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 87ee87cc6..58b3cc1ac 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 \ @@ -1035,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. From 72308d3f1cb460c4d3744eacff5c10e22ec75585 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:57:09 -0400 Subject: [PATCH 04/10] fix(ipc): allow one dictionary batch per ID when writing files On main, a file could contain exactly one DictionaryBatch, so a schema with more than one dictionary-encoded field couldn't be written, and the footer schema couldn't resolve its dictionary IDs when finalizing the file. Record the dictionary encodings of the footer schema and track the dictionary IDs written to the file: a second batch for the same ID returns EINVAL because IPC files don't support dictionary replacement (deltas remain unsupported when writing files). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 37 +++++++++++---- src/nanoarrow/ipc/writer_test.cc | 81 +++++++++++++++++++++++++++----- src/nanoarrow/nanoarrow_ipc.h | 4 ++ 3 files changed, 102 insertions(+), 20 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 65c0e449e..091236be1 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); } @@ -279,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; @@ -338,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)); @@ -370,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; diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index fba040fa8..b13458557 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -19,6 +19,12 @@ #include +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) +#include +#include +#include +#endif + #include "nanoarrow/nanoarrow_ipc.hpp" TEST(NanoarrowIpcWriter, OutputStreamBuffer) { @@ -277,28 +283,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); + 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. @@ -416,6 +426,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 58b3cc1ac..af903bc96 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1180,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( From 24b86176bdb829dec8e1f9ff4e912b240df44482 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:57:50 -0400 Subject: [PATCH 05/10] fix(ipc): write nested dictionaries before their parent When dictionary values contain dictionary-encoded children, a reader needs the nested dictionaries before it can decode the DictionaryBatch that references them. Keep assigning IDs in schema order but write each dictionary after the dictionaries nested in its values. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 17 ++-- src/nanoarrow/ipc/writer_test.cc | 128 +++++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+), 8 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 091236be1..6db55da5b 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -403,18 +403,17 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( // 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. +// A dictionary is written after any dictionaries nested in its values so that they +// are available when a reader decodes it. 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. static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view, 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++) { @@ -425,6 +424,8 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( if (array_view->dictionary != NULL) { NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView( writer, array_view->dictionary, next_id, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( + writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error)); } return NANOARROW_OK; diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index b13458557..cc675eb58 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -340,6 +340,134 @@ static void MakeDictionaryStructArray(struct ArrowArray* array, ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); } +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 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. From c673c8b08f2bfb774236a11027a462f2fb0f9e86 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 12:01:43 -0400 Subject: [PATCH 06/10] feat(ipc): don't repeat shared dictionaries when writing array streams ArrowIpcWriterWriteArrayStream() wrote every dictionary before every record batch. Skip a dictionary when it has the same length, offset, null count, and buffer addresses (recursively) as the corresponding dictionary of the previous array, which is the case when a producer shares one dictionary among the arrays of a stream (e.g., the nanoarrow IPC reader). Comparing addresses is only safe while both arrays are alive, so the writer keeps the previous array (if it has dictionaries) until the next one is written. ArrowIpcWriterWriteDictionaryBatch() only sees a borrowed view that it can't keep alive and always writes the dictionary. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 93 ++++++-- src/nanoarrow/ipc/writer_test.cc | 398 ++++++++++++++++++++++++++++++- src/nanoarrow/nanoarrow_ipc.h | 9 + 3 files changed, 480 insertions(+), 20 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 6db55da5b..e4f58a932 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -400,32 +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. // A dictionary is written after any dictionaries nested in its values so that they -// are available when a reader decodes it. 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. -static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView( +// 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) { (*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)); - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( - writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, 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; @@ -433,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)); @@ -448,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 @@ -466,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); @@ -480,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 cc675eb58..f542c34ce 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -18,6 +18,10 @@ #include #include +#include + +#include +#include #if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) #include @@ -313,7 +317,8 @@ TEST(NanoarrowIpcWriter, WriteDictionaryBatch) { // Build a struct array with a single dictionary-encoded (int32 -> utf8) child. static void MakeDictionaryStructArray(struct ArrowArray* array, - struct ArrowSchema* schema) { + struct ArrowSchema* schema, + const char* value1 = "bar") { ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), @@ -330,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); @@ -340,6 +345,42 @@ 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; @@ -468,6 +509,359 @@ TEST(NanoarrowIpcWriter, WritesNestedDictionariesDependencyFirst) { 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. diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index af903bc96..d45cc1bc8 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1192,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, From 3e4046aed5b42a0a375a84ff681d264818789f0d Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 12:04:16 -0400 Subject: [PATCH 07/10] test(ipc): roundtrip dictionary integration files through the writer Write the integration files back using ArrowIpcWriterWriteArrayStream(), which writes dictionaries (writing the arrays with ArrowIpcWriterWriteArrayView() alone does not), and enable the write roundtrip for the dictionary, nested dictionary, shared dictionary, and extension (whose dict-extension column is dictionary-encoded) files. Because the reader shares dictionaries among the batches it decodes, these roundtrips also exercise skipping unchanged dictionaries. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/files_test.cc | 48 ++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 21 deletions(-) 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") From 2b00267083444df1234dc5278a4d04631a0efb8e Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Fri, 16 Oct 2026 12:00:47 -0400 Subject: [PATCH 08/10] Don't compare view arrays' variadic sizes buffers; return ENOTSUP for a repeated file dictionary ID Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 27 ++++++++++++++++++++------- src/nanoarrow/ipc/writer_test.cc | 4 ++-- src/nanoarrow/nanoarrow_ipc.h | 10 ++++++---- 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index e4f58a932..8062c6cf6 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -361,7 +361,7 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( ArrowErrorSet( error, "Arrow IPC files do not support replacement of dictionary ID %" PRId64, dictionary_id); - return EINVAL; + return ENOTSUP; } } } @@ -404,8 +404,9 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( // 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, +// 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 || @@ -414,20 +415,31 @@ static int ArrowIpcWriterArraysShareBuffers(const struct ArrowArray* lhs, return 0; } - for (int64_t i = 0; i < lhs->n_buffers; i++) { + // 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(lhs->children[i], rhs->children[i])) { + if (!ArrowIpcWriterArraysShareBuffers(array_view->children[i], lhs->children[i], + rhs->children[i])) { return 0; } } return lhs->dictionary == NULL || - ArrowIpcWriterArraysShareBuffers(lhs->dictionary, rhs->dictionary); + ArrowIpcWriterArraysShareBuffers(array_view->dictionary, lhs->dictionary, + rhs->dictionary); } // Walk the array in the same depth-first order the schema encoder uses to assign @@ -461,7 +473,8 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArray( error)); if (previous_dictionary == NULL || - !ArrowIpcWriterArraysShareBuffers(array->dictionary, previous_dictionary)) { + !ArrowIpcWriterArraysShareBuffers(array_view->dictionary, array->dictionary, + previous_dictionary)) { NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch( writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error)); } diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index f542c34ce..d46eec2bf 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -291,7 +291,7 @@ TEST(NanoarrowIpcWriter, WriteDictionaryBatch) { int64_t bytes_written = p2->bytes_written; EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(), /*dictionary_id=*/0, /*is_delta=*/0, values_view.get(), &error), - EINVAL); + ENOTSUP); EXPECT_STREQ(error.message, "Arrow IPC files do not support replacement of dictionary ID 0"); EXPECT_EQ(p2->bytes_written, bytes_written); @@ -685,7 +685,7 @@ TEST(NanoarrowIpcWriter, RejectsChangedDictionaryInFile) { nanoarrow::UniqueBuffer output; EXPECT_EQ(WriteArrays(schema.get(), &arrays, output.get(), /*as_file=*/true, &error), - EINVAL); + ENOTSUP); EXPECT_STREQ(error.message, "Arrow IPC files do not support replacement of dictionary ID 0"); } diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index d45cc1bc8..73760182e 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1182,8 +1182,8 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// 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). +/// 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( @@ -1196,10 +1196,12 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( /// 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 +/// 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 EINVAL. +/// 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. From 57c4810478bdd4c2007a47024aff2f7877c9a17a Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 19 Oct 2026 15:25:21 -0400 Subject: [PATCH 09/10] feat(ipc): Skip unchanged dictionaries in ArrowIpcWriterWriteArrayViewWithDictionaries ArrowIpcWriterWriteArrayViewWithDictionaries() wrote every dictionary before every RecordBatch: its input is borrowed, so it couldn't compare a dictionary with the previous batch's by address as ArrowIpcWriterWriteArrayStream() does. The writer now keeps a copy of each dictionary it writes and compares the next batch's values with it, writing only dictionaries that changed (and those whose values contain a nested dictionary that was written). In a file, batches whose dictionaries are unchanged can now be written this way. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 313 ++++++++++++++++++++++++++++++- src/nanoarrow/ipc/writer_test.cc | 301 ++++++++++++++++++++++++++++- src/nanoarrow/nanoarrow_ipc.h | 21 ++- 3 files changed, 618 insertions(+), 17 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index fca3f308a..ec73e9345 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -189,8 +189,51 @@ struct ArrowIpcWriterPrivate { // DictionaryBatch messages written ahead of it) struct ArrowBuffer pending_message_metadata; int has_pending_message_metadata; + // Copies of the dictionaries last written by + // ArrowIpcWriterWriteArrayViewWithDictionaries() (struct ArrowIpcWriterViewDictionary) + struct ArrowBuffer view_dictionaries; }; +// The values of a dictionary as last written by +// ArrowIpcWriterWriteArrayViewWithDictionaries(). Its input is borrowed, so the +// values are copied to compare them with the next batch's. +struct ArrowIpcWriterViewDictionary { + int64_t dictionary_id; + struct ArrowArray values; +}; + +static void ArrowIpcWriterForgetViewDictionary(struct ArrowIpcWriterPrivate* private, + int64_t dictionary_id) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + if (entries[i].dictionary_id == dictionary_id) { + ArrowArrayRelease(&entries[i].values); + entries[i] = entries[n_entries - 1]; + NANOARROW_ASSERT_OK( + ArrowBufferResize(&private->view_dictionaries, + private->view_dictionaries.size_bytes - + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary), + /*shrink_to_fit=*/0)); + return; + } + } +} + +static void ArrowIpcWriterResetViewDictionaries(struct ArrowIpcWriterPrivate* private) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + ArrowArrayRelease(&entries[i].values); + } + NANOARROW_ASSERT_OK( + ArrowBufferResize(&private->view_dictionaries, 0, /*shrink_to_fit=*/0)); +} + ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, struct ArrowIpcOutputStream* output_stream) { NANOARROW_DCHECK(writer != NULL && output_stream != NULL); @@ -213,6 +256,7 @@ ArrowErrorCode ArrowIpcWriterInit(struct ArrowIpcWriter* writer, ArrowBufferInit(&private->file_dictionary_ids); ArrowBufferInit(&private->pending_message_metadata); private->has_pending_message_metadata = 0; + ArrowBufferInit(&private->view_dictionaries); writer->private_data = private; return NANOARROW_OK; @@ -233,6 +277,8 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { ArrowIpcFooterReset(&private->footer); ArrowBufferReset(&private->file_dictionary_ids); ArrowBufferReset(&private->pending_message_metadata); + ArrowIpcWriterResetViewDictionaries(private); + ArrowBufferReset(&private->view_dictionaries); ArrowFree(private); } @@ -299,6 +345,8 @@ ArrowErrorCode ArrowIpcWriterWriteSchema(struct ArrowIpcWriter* writer, // A dictionary may be replaced in a stream (e.g., when the dictionaries of // ArrowIpcWriterWriteArrayStream() input change between batches) but not in a file ArrowIpcEncoderSetDictionaryReplacement(&private->encoder, !private->writing_file); + // Dictionary IDs are assigned again for the new schema + ArrowIpcWriterResetViewDictionaries(private); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0)); NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSchema(&private->encoder, in, error)); @@ -372,10 +420,9 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* writer, return NANOARROW_OK; } -ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( +static ArrowErrorCode ArrowIpcWriterWriteDictionaryBatchImpl( struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta, const struct ArrowArrayView* values_view, struct ArrowError* error) { - NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && values_view != NULL); struct ArrowIpcWriterPrivate* private = (struct ArrowIpcWriterPrivate*)writer->private_data; @@ -432,6 +479,20 @@ ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( return NANOARROW_OK; } +ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch( + struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta, + const struct ArrowArrayView* values_view, struct ArrowError* error) { + NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && values_view != NULL); + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + + // ArrowIpcWriterWriteArrayViewWithDictionaries() must not skip this ID's next + // dictionary because it matches the values it wrote before this one + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + return ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, is_delta, + values_view, error); +} + // Check whether two arrays are guaranteed to contain identical values because they // have the same shape and point to the same buffers (e.g., a dictionary shared by // several arrays in a stream). This is only meaningful while both arrays are alive: @@ -516,15 +577,255 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArray( return NANOARROW_OK; } +// Initialize out as an array view with the same types as src +static ArrowErrorCode ArrowIpcWriterArrayViewInitLike(struct ArrowArrayView* out, + const struct ArrowArrayView* src) { + ArrowArrayViewInitFromType(out, src->storage_type); + out->layout = src->layout; + NANOARROW_RETURN_NOT_OK(ArrowArrayViewAllocateChildren(out, src->n_children)); + for (int64_t i = 0; i < src->n_children; i++) { + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterArrayViewInitLike(out->children[i], src->children[i])); + } + + if (src->dictionary != NULL) { + NANOARROW_RETURN_NOT_OK(ArrowArrayViewAllocateDictionary(out)); + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterArrayViewInitLike(out->dictionary, src->dictionary)); + } + + return NANOARROW_OK; +} + +// Producers need not initialize the padding bits of a bitmap's last byte. Clear them +// so that they don't make two copies of the same values compare unequal. +static void ArrowIpcWriterClearBitmapPadding(struct ArrowArray* array, + const struct ArrowArrayView* array_view) { + int64_t remainder = array->length % 8; + if (remainder != 0) { + uint8_t mask = (uint8_t)((1U << remainder) - 1U); + for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) { + if (array_view->layout.element_size_bits[i] == 1) { + struct ArrowBuffer* buffer = ArrowArrayBuffer(array, i); + if (buffer->size_bytes > 0) { + buffer->data[buffer->size_bytes - 1] &= mask; + } + } + } + } + + for (int64_t i = 0; i < array->n_children; i++) { + ArrowIpcWriterClearBitmapPadding(array->children[i], array_view->children[i]); + } +} + +// Copy the values of a dictionary into out, in a form that only depends on the values: +// zero offset, buffers sized to the values, and cleared bitmap padding. A dictionary +// nested in the values is not copied: it is written (and compared) separately. +static ArrowErrorCode ArrowIpcWriterCopyDictionaryValues( + const struct ArrowArrayView* values_view, struct ArrowArray* out, + struct ArrowError* error) { + NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromArrayView(out, values_view, error)); + ArrowErrorCode result = ArrowArrayStartAppending(out); + if (result == NANOARROW_OK) { + result = ArrowArrayReserve(out, values_view->length); + } + if (result == NANOARROW_OK) { + result = ArrowArrayAppendStorageFromArrayView(out, values_view, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayFinishBuildingDefault(out, error); + } + if (result != NANOARROW_OK) { + ArrowArrayRelease(out); + return result; + } + + ArrowIpcWriterClearBitmapPadding(out, values_view); + return NANOARROW_OK; +} + +// Compare two copies made by ArrowIpcWriterCopyDictionaryValues() from views of +// the same type as values_view +static ArrowErrorCode ArrowIpcWriterDictionaryValuesEqual( + const struct ArrowArrayView* values_view, const struct ArrowArray* lhs, + const struct ArrowArray* rhs, int* out, struct ArrowError* error) { + struct ArrowArrayView lhs_view; + struct ArrowArrayView rhs_view; + ArrowArrayViewInitFromType(&lhs_view, NANOARROW_TYPE_UNINITIALIZED); + ArrowArrayViewInitFromType(&rhs_view, NANOARROW_TYPE_UNINITIALIZED); + + ArrowErrorCode result = ArrowIpcWriterArrayViewInitLike(&lhs_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowIpcWriterArrayViewInitLike(&rhs_view, values_view); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&lhs_view, lhs, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&rhs_view, rhs, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewCompare(&lhs_view, &rhs_view, NANOARROW_COMPARE_IDENTICAL, out, + NULL); + } + + ArrowArrayViewReset(&lhs_view); + ArrowArrayViewReset(&rhs_view); + return result; +} + +static struct ArrowIpcWriterViewDictionary* ArrowIpcWriterFindViewDictionary( + struct ArrowIpcWriterPrivate* private, int64_t dictionary_id) { + struct ArrowIpcWriterViewDictionary* entries = + (struct ArrowIpcWriterViewDictionary*)private->view_dictionaries.data; + int64_t n_entries = private->view_dictionaries.size_bytes / + (int64_t)sizeof(struct ArrowIpcWriterViewDictionary); + for (int64_t i = 0; i < n_entries; i++) { + if (entries[i].dictionary_id == dictionary_id) { + return &entries[i]; + } + } + return NULL; +} + +// Check whether values_view is physically identical to copy (made by +// ArrowIpcWriterCopyDictionaryValues() from a view of the same type). If it is, the +// values are equal; if not, they may still be (e.g., if values_view has an offset). +static ArrowErrorCode ArrowIpcWriterViewMatchesCopy( + const struct ArrowArrayView* values_view, const struct ArrowArray* copy, int* out, + struct ArrowError* error) { + struct ArrowArrayView copy_view; + ArrowArrayViewInitFromType(©_view, NANOARROW_TYPE_UNINITIALIZED); + ArrowErrorCode result = ArrowIpcWriterArrayViewInitLike(©_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(©_view, copy, error); + } + if (result == NANOARROW_OK) { + result = ArrowArrayViewCompare(values_view, ©_view, NANOARROW_COMPARE_IDENTICAL, + out, NULL); + } + ArrowArrayViewReset(©_view); + return result; +} + +// Write the dictionary values_view for dictionary_id unless its values are identical +// to the ones this function last wrote for that ID (and force is not set). Sets +// *written to whether a DictionaryBatch was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( + struct ArrowIpcWriter* writer, int64_t dictionary_id, + const struct ArrowArrayView* values_view, int force, int* written, + struct ArrowError* error) { + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + + struct ArrowIpcWriterViewDictionary* entry = + ArrowIpcWriterFindViewDictionary(private, dictionary_id); + + // The common case: the values are physically identical to the copy (e.g., because + // the producer builds each batch's dictionary the same way) + if (entry != NULL && !force) { + int identical = 0; + NANOARROW_RETURN_NOT_OK( + ArrowIpcWriterViewMatchesCopy(values_view, &entry->values, &identical, error)); + if (identical) { + *written = 0; + return NANOARROW_OK; + } + } + + // Otherwise, compare a copy of the values made the same way (zero offset, cleared + // padding bits, and computed null count) + struct ArrowArray values; + if (ArrowIpcWriterCopyDictionaryValues(values_view, &values, error) != NANOARROW_OK) { + // These values can't be copied (e.g., a dense union), so they can't be compared + // with the next batch's either: write them every time + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + *written = 1; + return ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/0, values_view, error); + } + + ArrowErrorCode result = NANOARROW_OK; + if (entry != NULL && !force) { + int equal = 0; + result = ArrowIpcWriterDictionaryValuesEqual(values_view, &entry->values, &values, + &equal, error); + if (result == NANOARROW_OK && equal) { + ArrowArrayRelease(&values); + *written = 0; + return NANOARROW_OK; + } + } + + if (result == NANOARROW_OK) { + result = ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/0, values_view, error); + } + + if (result == NANOARROW_OK && entry == NULL) { + struct ArrowIpcWriterViewDictionary new_entry = {.dictionary_id = dictionary_id}; + ArrowArrayMove(&values, &new_entry.values); + result = + ArrowBufferAppend(&private->view_dictionaries, &new_entry, sizeof(new_entry)); + if (result != NANOARROW_OK) { + ArrowArrayRelease(&new_entry.values); + } + } else if (result == NANOARROW_OK) { + ArrowArrayRelease(&entry->values); + ArrowArrayMove(&values, &entry->values); + } else { + ArrowArrayRelease(&values); + } + + *written = result == NANOARROW_OK; + return result; +} + +// Write the dictionaries of a borrowed array view in the same order as +// ArrowIpcWriterWriteDictionariesForArray(), skipping those that are identical to +// the ones last written for their IDs. A dictionary whose values contain a dictionary +// that was written is written again even if its own values are unchanged: a reader +// resolves a DictionaryBatch against the nested dictionaries current at that point. +// Sets *written if any dictionary of array_view was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionaries( + struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view, + int64_t* next_id, int* written, struct ArrowError* error) { + int64_t dictionary_id = *next_id; + if (array_view->dictionary != NULL) { + (*next_id)++; + } + + for (int64_t i = 0; i < array_view->n_children; i++) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, array_view->children[i], next_id, written, error)); + } + + if (array_view->dictionary != NULL) { + int nested_written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, array_view->dictionary, next_id, &nested_written, error)); + + int dictionary_written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionary( + writer, dictionary_id, array_view->dictionary, /*force=*/nested_written, + &dictionary_written, error)); + if (nested_written || dictionary_written) { + *written = 1; + } + } + + return NANOARROW_OK; +} + ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( struct ArrowIpcWriter* writer, const struct ArrowArrayView* in, struct ArrowError* error) { NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && in != NULL); int64_t next_dictionary_id = 0; - // in is borrowed, so its dictionaries can't be kept alive to compare against the - // next batch's: write every one of them - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArray( - writer, in, /*array=*/NULL, /*previous=*/NULL, &next_dictionary_id, error)); + int written = 0; + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaries( + writer, in, &next_dictionary_id, &written, error)); return ArrowIpcWriterWriteArrayView(writer, in, error); } diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index a67cfb718..10077ad92 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -1480,9 +1480,9 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) { } TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesInFile) { - // A dictionary ID may only be written once in a file: writing one batch with its - // dictionaries succeeds, and a second is rejected rather than writing an invalid - // file + // A dictionary ID may only be written once in a file: batches whose dictionaries + // are unchanged can be written, and a changed dictionary is rejected rather than + // writing an invalid file struct ArrowError error; nanoarrow::UniqueSchema schema; @@ -1507,9 +1507,304 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesInFile) { &error), NANOARROW_OK) << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + + nanoarrow::UniqueSchema changed_schema; + nanoarrow::UniqueArray changed; + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(changed.get(), changed_schema.get(), "baz")); + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), changed.get(), &error), NANOARROW_OK) + << error.message; EXPECT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), &error), ENOTSUP); + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK) + << error.message; + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchFileReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + EXPECT_EQ(maybe_arrow_reader.ValueUnsafe()->num_record_batches(), 2); +#endif +} + +// Write each array with ArrowIpcWriterWriteArrayViewWithDictionaries() to a stream +static void WriteArrayViews(struct ArrowSchema* schema, + const std::vector& arrays, + struct ArrowBuffer* output) { + struct ArrowError error; + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema, &error), NANOARROW_OK) + << error.message; + + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema, &error), NANOARROW_OK) + << error.message; + for (struct ArrowArray* array : arrays) { + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + } + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesSkipsUnchanged) { + // Separately built arrays with the same dictionary values don't share any buffers, + // but only a changed dictionary is written again + std::vector schemas(4); + std::vector arrays(4); + const char* values[] = {"bar", "bar", "baz", "baz"}; + for (size_t i = 0; i < arrays.size(); i++) { + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(arrays[i].get(), schemas[i].get(), values[i])); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews( + schemas[0].get(), + {arrays[0].get(), arrays[1].get(), arrays[2].get(), arrays[3].get()}, + output.get())); + EXPECT_EQ(DecodeMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchStreamReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + for (const char* expected : values) { + std::shared_ptr arrow_batch; + ASSERT_TRUE(maybe_arrow_reader.ValueUnsafe()->ReadNext(&arrow_batch).ok()); + ASSERT_NE(arrow_batch, nullptr); + auto arrow_dictionary = + std::static_pointer_cast(arrow_batch->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + EXPECT_EQ(arrow_values->GetString(1), expected); + } +#endif + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 4); + for (size_t i = 0; i < roundtrip.size(); i++) { + EXPECT_EQ(StringValues(roundtrip[i]->children[0]->dictionary), + (std::vector{"foo", values[i]})); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesComparesValues) { + // Dictionaries are compared by value: a slice of a longer dictionary, an unknown null + // count, and garbage in a bitmap's padding bits don't make an unchanged dictionary + // look changed + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ASSERT_EQ( + ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_INT64), + NANOARROW_OK); + + // Values [1, null, 3], built as is (first) or as a slice of [0, 1, null, 3, 4] + std::vector arrays(2); + for (size_t i = 0; i < arrays.size(); i++) { + struct ArrowArray* array = arrays[i].get(); + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema.get(), nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendInt(values, 0), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 3), NANOARROW_OK); + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendInt(values, 4), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 2), NANOARROW_OK); + array->length = 2; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, &error), NANOARROW_OK) + << error.message; + if (i == 1) { + // Set the padding bits of the validity bitmap, slice to [1, null, 3], and leave + // the null count unknown + ArrowArrayBuffer(values, 0)->data[0] |= 0xe0; + values->offset = 1; + values->length = 3; + values->null_count = -1; + } + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schema.get(), {arrays[0].get(), arrays[1].get()}, output.get())); + EXPECT_EQ(DecodeMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNested) { + // A changed nested dictionary is written, and so is the dictionary whose values + // refer to it, even though that dictionary's own values are unchanged + const std::vector inner_values = {"bar", "bar", "baz"}; + std::vector schemas(inner_values.size()); + std::vector arrays(inner_values.size()); + for (size_t i = 0; i < inner_values.size(); i++) { + ASSERT_NO_FATAL_FAILURE(MakeNestedDictionaryStructArray( + arrays[i].get(), schemas[i].get(), inner_values[i].c_str())); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schemas[0].get(), + {arrays[0].get(), arrays[1].get(), arrays[2].get()}, output.get())); + EXPECT_EQ(DecodeDictionaryIds(output.get()), (std::vector{1, 0, 1, 0})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 3); + for (size_t i = 0; i < roundtrip.size(); i++) { + struct ArrowArray* outer_dictionary = roundtrip[i]->children[0]->dictionary; + EXPECT_EQ(StringValues(outer_dictionary->children[0]->dictionary), + (std::vector{"foo", inner_values[i]})); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesAfterWriteDictionaryBatch) { + // A dictionary written with ArrowIpcWriterWriteDictionaryBatch() replaces the one + // ArrowIpcWriterWriteArrayViewWithDictionaries() wrote, so the next batch's + // dictionary is written even though it is identical to that one + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(array.get(), schema.get())); + nanoarrow::UniqueSchema other_schema; + nanoarrow::UniqueArray other; + ASSERT_NO_FATAL_FAILURE( + MakeDictionaryStructArray(other.get(), other_schema.get(), "baz")); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView other_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(other_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(other_view.get(), other.get(), &error), NANOARROW_OK) + << error.message; + + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcWriterWriteDictionaryBatch(writer.get(), 0, /*is_delta=*/0, + other_view->children[0]->dictionary, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK); + + EXPECT_EQ(DecodeMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); + EXPECT_EQ(StringValues(roundtrip[1]->children[0]->dictionary), + (std::vector{"foo", "bar"})); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesUncopyableValues) { + // Dense union values can't be copied to compare them with the next batch's, so they + // are written for every batch + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ArrowSchemaInit(schema->children[0]->dictionary); + ASSERT_EQ(ArrowSchemaSetTypeUnion(schema->children[0]->dictionary, + NANOARROW_TYPE_DENSE_UNION, 1), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema->children[0]->dictionary->children[0], + NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0]->dictionary->children[0], "i"), + NANOARROW_OK); + + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + for (int32_t value : {5, 6}) { + ASSERT_EQ(ArrowArrayAppendInt(values->children[0], value), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishUnionElement(values, 0), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 1), NANOARROW_OK); + array->length = 1; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), &error), NANOARROW_OK) + << error.message; + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE( + WriteArrayViews(schema.get(), {array.get(), array.get()}, output.get())); + EXPECT_EQ(DecodeMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); } namespace { diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index d10f2fc18..1fe9b3620 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1211,12 +1211,16 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// assigned in the same order as when the schema was written, and nested dictionaries /// are written before the dictionaries that refer to them. /// -/// Because in is borrowed, its dictionaries are not compared with those of earlier -/// batches (as ArrowIpcWriterWriteArrayStream() does to skip unchanged ones): each -/// dictionary is written in full before every RecordBatch, replacing the previous -/// one. In file mode, where a dictionary ID may only be written once, this means only -/// one batch with dictionaries can be written this way. in must not be NULL; use -/// ArrowIpcWriterWriteArrayView() to write EOS. +/// A dictionary whose values are identical to the ones this function last wrote for +/// its ID is not written again. Because in is borrowed, the writer keeps a copy of +/// each dictionary it writes to compare with the next batch's (so the comparison costs +/// a copy of each dictionary per batch); values that can't be copied (e.g., a dense +/// union) are written every time. A changed dictionary is written in full, replacing +/// the previous one, which a file does not support (ENOTSUP). A dictionary whose +/// values contain a dictionary that was written is written again too. +/// ArrowIpcWriterWriteDictionaryBatch() and ArrowIpcWriterWriteSchema() clear the +/// copies they affect. in must not be NULL; use ArrowIpcWriterWriteArrayView() to +/// write EOS. /// /// Errors are propagated from the underlying encoder and output byte stream. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( @@ -1228,8 +1232,9 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries( /// dictionary_id must match the id assigned to the dictionary-encoded field in the /// schema. is_delta selects DictionaryBatch.isDelta. values_view must not itself be /// dictionary-encoded. The writer does not check that a schema was already written. -/// values_view is always written: the writer only sees a borrowed view and cannot -/// check whether it is identical to a dictionary that was already written. When +/// values_view is always written: unlike ArrowIpcWriterWriteArrayViewWithDictionaries(), +/// this does not check whether it is identical to a dictionary that was already +/// written. When /// writing a file, returns ENOTSUP for a delta or for a dictionary_id that was already /// written (IPC files do not support dictionary replacement). /// From cf1f753522631e82fbb79fb6b6d37c76eec901f1 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 19 Oct 2026 15:37:19 -0400 Subject: [PATCH 10/10] feat(ipc): Write dictionary deltas from ArrowIpcWriterWriteArrayViewWithDictionaries When a batch's dictionary starts with the values last written for its ID, the writer now writes only the appended values, as a delta DictionaryBatch, instead of replacing the whole dictionary. This only applies to streams: a file does not support deltas. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/nanoarrow/ipc/writer.c | 113 ++++++++++++++-- src/nanoarrow/ipc/writer_test.cc | 215 +++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 8 +- 3 files changed, 323 insertions(+), 13 deletions(-) diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index ec73e9345..b6e9d7921 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -619,19 +619,25 @@ static void ArrowIpcWriterClearBitmapPadding(struct ArrowArray* array, } } -// Copy the values of a dictionary into out, in a form that only depends on the values: -// zero offset, buffers sized to the values, and cleared bitmap padding. A dictionary -// nested in the values is not copied: it is written (and compared) separately. +// Copy length values of a dictionary, starting at offset, into out, in a form that +// only depends on the values: zero offset, buffers sized to the values, and cleared +// bitmap padding. A dictionary nested in the values is not copied: it is written (and +// compared) separately. static ArrowErrorCode ArrowIpcWriterCopyDictionaryValues( - const struct ArrowArrayView* values_view, struct ArrowArray* out, - struct ArrowError* error) { + const struct ArrowArrayView* values_view, int64_t offset, int64_t length, + struct ArrowArray* out, struct ArrowError* error) { + struct ArrowArrayView slice = *values_view; + slice.offset += offset; + slice.length = length; + slice.null_count = -1; + NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromArrayView(out, values_view, error)); ArrowErrorCode result = ArrowArrayStartAppending(out); if (result == NANOARROW_OK) { - result = ArrowArrayReserve(out, values_view->length); + result = ArrowArrayReserve(out, length); } if (result == NANOARROW_OK) { - result = ArrowArrayAppendStorageFromArrayView(out, values_view, error); + result = ArrowArrayAppendStorageFromArrayView(out, &slice, error); } if (result == NANOARROW_OK) { result = ArrowArrayFinishBuildingDefault(out, error); @@ -709,9 +715,84 @@ static ArrowErrorCode ArrowIpcWriterViewMatchesCopy( return result; } +// If values_view starts with the values last written for entry, write the values +// appended to them as a delta and add them to entry's copy. Sets *written to whether a +// delta was written. +static ArrowErrorCode ArrowIpcWriterWriteViewDictionaryDelta( + struct ArrowIpcWriter* writer, struct ArrowIpcWriterViewDictionary* entry, + int64_t dictionary_id, const struct ArrowArrayView* values_view, int* written, + struct ArrowError* error) { + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + *written = 0; + + int64_t n_written = entry->values.length; + struct ArrowArray prefix; + if (ArrowIpcWriterCopyDictionaryValues(values_view, 0, n_written, &prefix, error) != + NANOARROW_OK) { + // Leave it to the caller, which writes values it can't copy in full + return NANOARROW_OK; + } + + int prefix_equal = 0; + ArrowErrorCode result = ArrowIpcWriterDictionaryValuesEqual( + values_view, &entry->values, &prefix, &prefix_equal, error); + if (result != NANOARROW_OK || !prefix_equal) { + ArrowArrayRelease(&prefix); + return result; + } + + struct ArrowArray delta; + result = ArrowIpcWriterCopyDictionaryValues( + values_view, n_written, values_view->length - n_written, &delta, error); + if (result != NANOARROW_OK) { + ArrowArrayRelease(&prefix); + return result; + } + + struct ArrowArrayView delta_view; + ArrowArrayViewInitFromType(&delta_view, NANOARROW_TYPE_UNINITIALIZED); + result = ArrowIpcWriterArrayViewInitLike(&delta_view, values_view); + if (result == NANOARROW_OK) { + result = ArrowArrayViewSetArray(&delta_view, &delta, error); + } + if (result == NANOARROW_OK) { + result = ArrowIpcWriterWriteDictionaryBatchImpl(writer, dictionary_id, + /*is_delta=*/1, &delta_view, error); + } + + if (result == NANOARROW_OK) { + *written = 1; + + // The values that were already written followed by the delta are a copy of + // values_view. If it can't be completed, forget the dictionary: the next batch's + // is then written in full. + ArrowErrorCode copy_result = + ArrowArrayAppendStorageFromArrayView(&prefix, &delta_view, error); + if (copy_result == NANOARROW_OK) { + copy_result = ArrowArrayFinishBuildingDefault(&prefix, error); + } + if (copy_result == NANOARROW_OK) { + ArrowIpcWriterClearBitmapPadding(&prefix, values_view); + ArrowArrayRelease(&entry->values); + ArrowArrayMove(&prefix, &entry->values); + } else { + ArrowIpcWriterForgetViewDictionary(private, dictionary_id); + } + } + + ArrowArrayViewReset(&delta_view); + ArrowArrayRelease(&delta); + if (prefix.release != NULL) { + ArrowArrayRelease(&prefix); + } + return result; +} + // Write the dictionary values_view for dictionary_id unless its values are identical -// to the ones this function last wrote for that ID (and force is not set). Sets -// *written to whether a DictionaryBatch was written. +// to the ones this function last wrote for that ID (and force is not set). In a stream, +// values that start with the ones last written are written as a delta of the values +// that were appended. Sets *written to whether a DictionaryBatch was written. static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( struct ArrowIpcWriter* writer, int64_t dictionary_id, const struct ArrowArrayView* values_view, int force, int* written, @@ -734,10 +815,22 @@ static ArrowErrorCode ArrowIpcWriterWriteViewDictionary( } } + // In a stream, values that start with the ones last written only need the appended + // values to be written, as a delta (a file does not support deltas) + if (entry != NULL && !force && !private->writing_file && + values_view->length > entry->values.length) { + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteViewDictionaryDelta( + writer, entry, dictionary_id, values_view, written, error)); + if (*written) { + return NANOARROW_OK; + } + } + // Otherwise, compare a copy of the values made the same way (zero offset, cleared // padding bits, and computed null count) struct ArrowArray values; - if (ArrowIpcWriterCopyDictionaryValues(values_view, &values, error) != NANOARROW_OK) { + if (ArrowIpcWriterCopyDictionaryValues(values_view, 0, values_view->length, &values, + error) != NANOARROW_OK) { // These values can't be copied (e.g., a dense union), so they can't be compared // with the next batch's either: write them every time ArrowIpcWriterForgetViewDictionary(private, dictionary_id); diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 10077ad92..35fa9487e 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -1807,6 +1807,221 @@ TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesUncopyableValues) { ASSERT_EQ(roundtrip.size(), 2); } +// The isDelta flag of each DictionaryBatch message of an encapsulated IPC stream +static std::vector DecodeDictionaryDeltaFlags(const struct ArrowBuffer* buffer) { + std::vector is_delta; + struct ArrowBufferView remaining; + remaining.data.as_uint8 = buffer->data; + remaining.size_bytes = buffer->size_bytes; + nanoarrow::ipc::UniqueDecoder decoder; + EXPECT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowError error; + + while (remaining.size_bytes > 0) { + int result = ArrowIpcDecoderVerifyHeader(decoder.get(), remaining, &error); + if (result == ENODATA) { + break; + } + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + + if (decoder->message_type == NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH) { + result = ArrowIpcDecoderDecodeHeader(decoder.get(), remaining, &error); + EXPECT_EQ(result, NANOARROW_OK) << error.message; + if (result != NANOARROW_OK) { + break; + } + is_delta.push_back(decoder->dictionary->is_delta); + } + + int64_t message_size = decoder->header_size_bytes + decoder->body_size_bytes; + remaining.data.as_uint8 += message_size; + remaining.size_bytes -= message_size; + } + return is_delta; +} + +// A struct> array of one row per dictionary value +static void MakeStringDictionaryStructArray(struct ArrowArray* array, + struct ArrowSchema* schema, + const std::vector& values) { + ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ASSERT_EQ( + ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_STRING), + NANOARROW_OK); + + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + for (size_t i = 0; i < values.size(); i++) { + ASSERT_EQ(ArrowArrayAppendString(array->children[0]->dictionary, + ArrowCharView(values[i].c_str())), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], static_cast(i)), + NANOARROW_OK); + } + array->length = static_cast(values.size()); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesDeltas) { + // A dictionary that only grew is written as a delta of its new values; any other + // change replaces it + const std::vector> batch_values = { + {"foo", "bar"}, {"foo", "bar", "baz"}, {"foo", "bar", "baz"}, {"foo", "qux"}}; + std::vector schemas(batch_values.size()); + std::vector arrays(batch_values.size()); + std::vector array_ptrs; + for (size_t i = 0; i < arrays.size(); i++) { + ASSERT_NO_FATAL_FAILURE(MakeStringDictionaryStructArray( + arrays[i].get(), schemas[i].get(), batch_values[i])); + array_ptrs.push_back(arrays[i].get()); + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews(schemas[0].get(), array_ptrs, output.get())); + EXPECT_EQ(DecodeMessageTypes(output.get()), + (std::vector{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH, + NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH})); + EXPECT_EQ(DecodeDictionaryDeltaFlags(output.get()), (std::vector{0, 1, 0})); + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + auto arrow_input = std::make_shared( + arrow::Buffer::Wrap(output->data, output->size_bytes)); + auto maybe_arrow_reader = arrow::ipc::RecordBatchStreamReader::Open(arrow_input); + ASSERT_TRUE(maybe_arrow_reader.ok()) << maybe_arrow_reader.status(); + for (const auto& expected : batch_values) { + std::shared_ptr arrow_batch; + auto status = maybe_arrow_reader.ValueUnsafe()->ReadNext(&arrow_batch); + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(arrow_batch, nullptr); + auto arrow_dictionary = + std::static_pointer_cast(arrow_batch->column(0)); + auto arrow_values = + std::static_pointer_cast(arrow_dictionary->dictionary()); + ASSERT_EQ(arrow_values->length(), static_cast(expected.size())); + for (size_t i = 0; i < expected.size(); i++) { + EXPECT_EQ(arrow_values->GetString(static_cast(i)), expected[i]); + } + } +#endif + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), batch_values.size()); + for (size_t i = 0; i < roundtrip.size(); i++) { + EXPECT_EQ(StringValues(roundtrip[i]->children[0]->dictionary), batch_values[i]); + } +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNullableDelta) { + // The appended values of a delta start partway through a byte of the validity + // bitmap, and the values that were already written still compare equal + struct ArrowError error; + std::vector schemas(2); + std::vector arrays(2); + for (size_t i = 0; i < arrays.size(); i++) { + struct ArrowSchema* schema = schemas[i].get(); + ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK); + ASSERT_EQ( + ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_INT64), + NANOARROW_OK); + + struct ArrowArray* array = arrays[i].get(); + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + struct ArrowArray* values = array->children[0]->dictionary; + // [1, null, 3] and then [1, null, 3, null, 5] + ASSERT_EQ(ArrowArrayAppendInt(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 3), NANOARROW_OK); + if (i == 1) { + ASSERT_EQ(ArrowArrayAppendNull(values, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(values, 5), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], 0), NANOARROW_OK); + array->length = 1; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, &error), NANOARROW_OK) + << error.message; + } + + nanoarrow::UniqueBuffer output; + ASSERT_NO_FATAL_FAILURE(WriteArrayViews( + schemas[0].get(), {arrays[0].get(), arrays[1].get()}, output.get())); + EXPECT_EQ(DecodeDictionaryDeltaFlags(output.get()), (std::vector{0, 1})); + + std::vector roundtrip; + ReadArrays(output.get(), &roundtrip); + ASSERT_EQ(roundtrip.size(), 2); + nanoarrow::UniqueArrayView view; + ArrowArrayViewInitFromType(view.get(), NANOARROW_TYPE_INT64); + ASSERT_EQ( + ArrowArrayViewSetArray(view.get(), roundtrip[1]->children[0]->dictionary, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(view->length, 5); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 0), 1); + EXPECT_TRUE(ArrowArrayViewIsNull(view.get(), 1)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 2), 3); + EXPECT_TRUE(ArrowArrayViewIsNull(view.get(), 3)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(view.get(), 4), 5); +} + +TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesNoDeltaInFile) { + // A file does not support deltas, so a dictionary that grew is a replacement there + struct ArrowError error; + std::vector schemas(2); + std::vector arrays(2); + ASSERT_NO_FATAL_FAILURE( + MakeStringDictionaryStructArray(arrays[0].get(), schemas[0].get(), {"foo", "bar"})); + ASSERT_NO_FATAL_FAILURE(MakeStringDictionaryStructArray( + arrays[1].get(), schemas[1].get(), {"foo", "bar", "baz"})); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schemas[0].get(), &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterStartFile(writer.get(), &error), NANOARROW_OK) << error.message; + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schemas[0].get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), arrays[0].get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), arrays[1].get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(), + &error), + ENOTSUP); +} + namespace { // Walk the messages of an encapsulated IPC stream, recording each message's type and diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 1fe9b3620..8c40daf62 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -1215,9 +1215,11 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter* /// its ID is not written again. Because in is borrowed, the writer keeps a copy of /// each dictionary it writes to compare with the next batch's (so the comparison costs /// a copy of each dictionary per batch); values that can't be copied (e.g., a dense -/// union) are written every time. A changed dictionary is written in full, replacing -/// the previous one, which a file does not support (ENOTSUP). A dictionary whose -/// values contain a dictionary that was written is written again too. +/// union) are written every time. In a stream, a dictionary that only grew (its values +/// start with the ones last written) is written as a delta of the appended values. +/// Otherwise, a changed dictionary is written in full, replacing the previous one, +/// which a file does not support (ENOTSUP). A dictionary whose +/// values contain a dictionary that was written is written again too, in full. /// ArrowIpcWriterWriteDictionaryBatch() and ArrowIpcWriterWriteSchema() clear the /// copies they affect. in must not be NULL; use ArrowIpcWriterWriteArrayView() to /// write EOS.