From 2925439d43d773e2ccdb05b301ae33396334bb86 Mon Sep 17 00:00:00 2001 From: ump45nose <52391318+ump45nose@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:10:22 +0800 Subject: [PATCH] fix(ipc): Use the decoder's decompressor for dictionary batches Compressed DictionaryBatch values were decompressed by the nested dictionary decoder's default serial decompressor, so a decompressor set with ArrowIpcDecoderSetDecompressor() was ignored for dictionaries. Pass the parent decoder's decompressor down when decoding the dictionary values. Closes #936. --- src/nanoarrow/ipc/decoder.c | 32 +++++++++---- src/nanoarrow/ipc/encoder_test.cc | 74 ++++++++++++++++++++++++++++++- 2 files changed, 96 insertions(+), 10 deletions(-) diff --git a/src/nanoarrow/ipc/decoder.c b/src/nanoarrow/ipc/decoder.c index 80e53ec8d..f9752ba07 100644 --- a/src/nanoarrow/ipc/decoder.c +++ b/src/nanoarrow/ipc/decoder.c @@ -2667,10 +2667,13 @@ static ArrowErrorCode ArrowIpcDecoderDecodeArrayInternal( return NANOARROW_OK; } +// If decompressor is NULL, the decoder's own decompressor is used (initializing the +// default decompressor if one has not been set). static ArrowErrorCode ArrowIpcDecoderDecodeArrayViewInternal( struct ArrowIpcDecoder* decoder, struct ArrowIpcBufferFactory factory, int64_t field_i, struct ArrowIpcDictionaries* dictionaries, - struct ArrowArrayView** out_view, struct ArrowError* error) { + struct ArrowIpcDecompressor* decompressor, struct ArrowArrayView** out_view, + struct ArrowError* error) { struct ArrowIpcDecoderPrivate* private_data = (struct ArrowIpcDecoderPrivate*)decoder->private_data; @@ -2700,8 +2703,11 @@ static ArrowErrorCode ArrowIpcDecoderDecodeArrayViewInternal( // If we are going to need a decompressor here, ensure the default one is // initialized. if (setter.src.codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { - NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderInitDecompressor(private_data)); - setter.factory.decompressor = &private_data->decompressor; + if (decompressor == NULL) { + NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderInitDecompressor(private_data)); + decompressor = &private_data->decompressor; + } + setter.factory.decompressor = decompressor; } // The flatbuffers FieldNode doesn't count the root struct so we have to loop over the @@ -2746,7 +2752,7 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderDecodeArrayViewWithDictionaries( } return ArrowIpcDecoderDecodeArrayViewInternal( - decoder, ArrowIpcBufferFactoryFromView(&body), i, dictionaries, out, error); + decoder, ArrowIpcBufferFactoryFromView(&body), i, dictionaries, NULL, out, error); } ArrowErrorCode ArrowIpcDecoderDecodeArrayView(struct ArrowIpcDecoder* decoder, @@ -2771,7 +2777,7 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderDecodeArrayWithDictionaries( struct ArrowArrayView* array_view; NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderDecodeArrayViewInternal( - decoder, ArrowIpcBufferFactoryFromView(&body), i, dictionaries, &array_view, + decoder, ArrowIpcBufferFactoryFromView(&body), i, dictionaries, NULL, &array_view, error)); NANOARROW_RETURN_NOT_OK(ArrowArrayViewValidate(array_view, validation_level, error)); @@ -2805,7 +2811,7 @@ ArrowErrorCode ArrowIpcDecoderDecodeArrayFromSharedWithDictionaries( enum ArrowValidationLevel validation_level, struct ArrowError* error) { struct ArrowArrayView* array_view; NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderDecodeArrayViewInternal( - decoder, ArrowIpcBufferFactoryFromShared(body), i, dictionaries, &array_view, + decoder, ArrowIpcBufferFactoryFromShared(body), i, dictionaries, NULL, &array_view, error)); NANOARROW_RETURN_NOT_OK(ArrowArrayViewValidate(array_view, validation_level, error)); @@ -2864,17 +2870,25 @@ static ArrowErrorCode ArrowIpcDecoderDecodeDictionaryInternal( (struct ArrowIpcDecoderPrivate*)dictionary->decoder.private_data; dictionary->decoder.message_type = NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH; dictionary->decoder.metadata_version = decoder->metadata_version; - // Transfer the endianness and compression settings so that buffers are byte-swapped - // and decompressed if needed (the nested decoder uses a default decompressor) + // Transfer the endianness so that buffers are byte-swapped if needed dictionary_decoder_private_data->endianness = private_data->endianness; // Validate the nested batch's node and buffer counts before accessing its vectors. NANOARROW_RETURN_NOT_OK( ArrowIpcDecoderDecodeRecordBatchHeader(&dictionary->decoder, record_batch, error)); dictionary_decoder_private_data->last_message = record_batch; + // Decompress the dictionary values with this decoder's decompressor (which may have + // been set by the caller with ArrowIpcDecoderSetDecompressor()) rather than with the + // nested decoder's default decompressor. + struct ArrowIpcDecompressor* decompressor = NULL; + if (dictionary->decoder.codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { + NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderInitDecompressor(private_data)); + decompressor = &private_data->decompressor; + } + struct ArrowArrayView* array_view; NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderDecodeArrayViewInternal( - &dictionary->decoder, factory, 0, dictionaries, &array_view, error)); + &dictionary->decoder, factory, 0, dictionaries, decompressor, &array_view, error)); NANOARROW_RETURN_NOT_OK(ArrowArrayViewValidate(array_view, validation_level, error)); diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index fa6b37b3a..a536e1146 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -1561,7 +1561,59 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedAllocationFailures) { } // DictionaryBatch bodies are compressed like RecordBatch bodies and must declare it -static void TestCompressedDictionaryBatch(enum ArrowIpcCompressionType codec) { +// A decompressor that forwards to the serial decompressor and counts how many +// buffers it was asked to decompress +struct CountingDecompressorPrivate { + struct ArrowIpcDecompressor wrapped; + int64_t* n_calls; +}; + +static ArrowErrorCode CountingDecompressorAdd( + struct ArrowIpcDecompressor* decompressor, + enum ArrowIpcCompressionType compression_type, struct ArrowBufferView src, + uint8_t* dst, int64_t dst_size, struct ArrowError* error) { + auto private_data = + reinterpret_cast(decompressor->private_data); + ++*private_data->n_calls; + return private_data->wrapped.decompress_add(&private_data->wrapped, compression_type, + src, dst, dst_size, error); +} + +static ArrowErrorCode CountingDecompressorWait(struct ArrowIpcDecompressor* decompressor, + int64_t timeout_ms, + struct ArrowError* error) { + auto private_data = + reinterpret_cast(decompressor->private_data); + return private_data->wrapped.decompress_wait(&private_data->wrapped, timeout_ms, error); +} + +static void CountingDecompressorRelease(struct ArrowIpcDecompressor* decompressor) { + auto private_data = + reinterpret_cast(decompressor->private_data); + private_data->wrapped.release(&private_data->wrapped); + delete private_data; + decompressor->release = nullptr; +} + +static ArrowErrorCode CountingDecompressor(struct ArrowIpcDecompressor* out, + int64_t* n_calls) { + auto private_data = new CountingDecompressorPrivate; + private_data->n_calls = n_calls; + ArrowErrorCode result = ArrowIpcSerialDecompressor(&private_data->wrapped); + if (result != NANOARROW_OK) { + delete private_data; + return result; + } + + out->decompress_add = &CountingDecompressorAdd; + out->decompress_wait = &CountingDecompressorWait; + out->release = &CountingDecompressorRelease; + out->private_data = private_data; + return NANOARROW_OK; +} + +static void TestCompressedDictionaryBatch(enum ArrowIpcCompressionType codec, + bool custom_decompressor = false) { struct ArrowError error; // A dictionary-encoded int32 -> utf8 field, which gets dictionary id 0 @@ -1638,6 +1690,12 @@ static void TestCompressedDictionaryBatch(enum ArrowIpcCompressionType codec) { ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); ASSERT_EQ(ArrowIpcDecoderSetEndianness(decoder.get(), ArrowIpcSystemEndianness()), NANOARROW_OK); + int64_t n_decompress_calls = 0; + if (custom_decompressor) { + struct ArrowIpcDecompressor decompressor; + ASSERT_EQ(CountingDecompressor(&decompressor, &n_decompress_calls), NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderSetDecompressor(decoder.get(), &decompressor), NANOARROW_OK); + } struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), NANOARROW_OK) @@ -1678,6 +1736,12 @@ static void TestCompressedDictionaryBatch(enum ArrowIpcCompressionType codec) { EXPECT_EQ(is_equal, 1) << error.message; } + // The dictionary values must be decompressed with the decompressor that was set + // on the decoder + if (custom_decompressor) { + EXPECT_GT(n_decompress_calls, 0); + } + ArrowIpcDictionariesReset(&dictionaries); ArrowIpcDictionaryEncodingsReset(&encodings); } @@ -1696,6 +1760,14 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedDictionaryBatchZstd) { TestCompressedDictionaryBatch(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); } +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedDictionaryBatchCustomDecompressor) { + if (ArrowIpcGetZstdCompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD"; + } + TestCompressedDictionaryBatch(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + /*custom_decompressor=*/true); +} + // Schemas encoded while a compressor is set declare the COMPRESSED_BODY feature TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaDeclaresCompression) { nanoarrow::ipc::UniqueEncoder encoder;