From a34ce843042edd9abc52794aca13a49b3fd76896 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Sun, 4 Oct 2026 11:42:50 -0400 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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")