Skip to content

fix(ipc): Use the decoder's decompressor for dictionary batches #994

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ump45nose wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(ipc): Use the decoder's decompressor for dictionary batches #994

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 23 additions & 9 deletions src/nanoarrow/ipc/decoder.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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));
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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));

Expand Down
74 changes: 73 additions & 1 deletion src/nanoarrow/ipc/encoder_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<CountingDecompressorPrivate*>(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<CountingDecompressorPrivate*>(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<CountingDecompressorPrivate*>(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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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);
}
Expand All @@ -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;
Expand Down