diff --git a/cpp/src/parquet/decoder.cc b/cpp/src/parquet/decoder.cc index c4d3fe5a8a5..e6afb178ba7 100644 --- a/cpp/src/parquet/decoder.cc +++ b/cpp/src/parquet/decoder.cc @@ -97,6 +97,20 @@ struct ArrowBinaryHelper { return Status::OK(); } + // Keep the common per-value path small when callers already know the remaining size. + ARROW_FORCE_INLINE Status AppendValue(const uint8_t* data, int32_t length, + int64_t estimated_remaining_data_length) { + DCHECK_GT(entries_remaining_, 0); + + if (ARROW_PREDICT_FALSE(!CanFit(length))) { + return AppendValueWithKnownSizeSlow(data, length, estimated_remaining_data_length); + } + chunk_space_remaining_ -= length; + --entries_remaining_; + builder_->UnsafeAppend(data, length); + return Status::OK(); + } + // If a new chunk is created and estimated_remaining_data_length is provided, // it will also reserve the estimated data length for this chunk. Status AppendValue(const uint8_t* data, int32_t length, @@ -133,6 +147,9 @@ struct ArrowBinaryHelper { } private: + ARROW_NOINLINE Status AppendValueWithKnownSizeSlow( + const uint8_t* data, int32_t length, int64_t estimated_remaining_data_length); + Status PushChunk() { ARROW_ASSIGN_OR_RAISE(auto chunk, acc_->builder->Finish()); acc_->chunks.push_back(std::move(chunk)); @@ -1024,8 +1041,7 @@ class DictDecoderImpl : public TypedDecoderImpl, public DictDecoder // memory use in most cases std::shared_ptr byte_array_offsets_; - // Reusable buffer for decoding dictionary indices to be appended to a - // BinaryDictionary32Builder + // Reusable buffer for decoding dictionary indices into Arrow builders. std::shared_ptr indices_scratch_space_; ::arrow::util::RleBitPackedDecoder idx_decoder_; @@ -1257,6 +1273,11 @@ void DictDecoderImpl::InsertDictionary(::arrow::ArrayBuilder* bui PARQUET_THROW_NOT_OK(binary_builder->InsertMemoValues(*arr)); } +// Keep defensive bitmap validation out of the hot dictionary decode function. +ARROW_NOINLINE void ValidateBinaryValidityBitmap(int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset); + class DictByteArrayDecoderImpl : public DictDecoderImpl { public: using BASE = DictDecoderImpl; @@ -1284,6 +1305,17 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl { /*valid_bits=*/nullptr, valid_bits_offset, out, &result)); } else { + switch (out->builder->type()->id()) { + case ::arrow::Type::BINARY: + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::LARGE_STRING: + ValidateBinaryValidityBitmap(num_values, null_count, valid_bits, + valid_bits_offset); + break; + default: + break; + } PARQUET_THROW_NOT_OK(DecodeArrowDense(num_values, null_count, valid_bits, valid_bits_offset, out, &result)); } @@ -1295,12 +1327,82 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl { int64_t valid_bits_offset, typename EncodingTraits::Accumulator* out, int* out_num_values) { + const auto* dict_values = dictionary_->data_as(); + const int values_to_decode = num_values - null_count; + + switch (out->builder->type()->id()) { + case ::arrow::Type::BINARY: + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::LARGE_STRING: { + if (values_to_decode > 0) { + RETURN_NOT_OK(indices_scratch_space_->TypedResize( + values_to_decode, /*shrink_to_fit=*/false)); + } + auto* decoded_indices = indices_scratch_space_->mutable_data_as(); + const int num_indices = idx_decoder_.GetBatch(decoded_indices, values_to_decode); + if (ARROW_PREDICT_FALSE(num_indices != values_to_decode)) { + return Status::Invalid( + "Invalid or truncated dictionary index stream: expected ", values_to_decode, + " indices but decoded ", num_indices); + } + + int64_t data_length = 0; + for (int i = 0; i < values_to_decode; ++i) { + const auto index = decoded_indices[i]; + RETURN_NOT_OK(IndexInBounds(index)); + if (ARROW_PREDICT_FALSE(AddWithOverflow( + data_length, static_cast(dict_values[index].len), + &data_length))) { + return Status::Invalid( + "excess expansion while decoding dictionary-encoded BYTE_ARRAY"); + } + } + + auto append_predecoded = [&](auto* helper) { + int values_decoded = 0; + int pos_indices = 0; + int64_t remaining_data_length = data_length; + + RETURN_NOT_OK(VisitBitRuns( + valid_bits, valid_bits_offset, num_values, + [&](int64_t position, int64_t length, bool valid) { + if (valid) { + for (int64_t i = 0; i < length; ++i) { + const auto& val = dict_values[decoded_indices[pos_indices++]]; + RETURN_NOT_OK(helper->AppendValue( + val.ptr, static_cast(val.len), remaining_data_length)); + remaining_data_length -= val.len; + } + values_decoded += static_cast(length); + } else { + for (int64_t i = 0; i < length; ++i) { + helper->UnsafeAppendNull(); + } + } + return Status::OK(); + })); + DCHECK_EQ(pos_indices, values_to_decode); + DCHECK_EQ(remaining_data_length, 0); + *out_num_values = values_decoded; + return Status::OK(); + }; + + return DispatchArrowBinaryHelper(out, num_values, data_length, + append_predecoded); + } + default: + // Binary-view builders don't benefit from reserving the dictionary values' + // total byte length, since short values are stored inline. Keep their + // existing bounded streaming decode path. Unsupported builder types are + // rejected by DispatchArrowBinaryHelper below before decoding any indices. + break; + } + constexpr int32_t kBufferSize = 1024; int32_t indices[kBufferSize]; - auto visit_binary_helper = [&](auto* helper) { - const auto* dict_values = dictionary_->data_as(); - const int values_to_decode = num_values - null_count; + auto append_streaming = [&](auto* helper) { int values_decoded = 0; int num_indices = 0; int pos_indices = 0; @@ -1309,7 +1411,7 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl { if (valid) { while (length > 0) { if (num_indices == pos_indices) { - // Refill indices buffer + // Refill the bounded indices buffer for binary-view builders. const auto max_batch_size = std::min(kBufferSize, values_to_decode - values_decoded); num_indices = idx_decoder_.GetBatch(indices, max_batch_size); @@ -1341,11 +1443,9 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl { *out_num_values = values_decoded; return Status::OK(); }; - // The `len_` in the ByteArrayDictDecoder is the total length of the - // RLE/Bit-pack encoded data size, so, we cannot use `len_` to reserve - // space for binary data. + return DispatchArrowBinaryHelper( - out, num_values, /*estimated_data_length=*/{}, visit_binary_helper); + out, num_values, /*estimated_data_length=*/{}, append_streaming); } template @@ -2372,6 +2472,33 @@ class ByteStreamSplitDecoder : public ByteStreamSplitDecoderBase::AppendValueWithKnownSizeSlow( + const uint8_t* data, int32_t length, int64_t estimated_remaining_data_length) { + RETURN_NOT_OK(PushChunk()); + // Validate that the first value fits in an empty chunk before UnsafeAppend below. + RETURN_NOT_OK(builder_->ReserveData(length)); + RETURN_NOT_OK(ReserveInitialChunkData(estimated_remaining_data_length)); + chunk_space_remaining_ -= length; + --entries_remaining_; + builder_->UnsafeAppend(data, length); + return Status::OK(); +} + +void ValidateBinaryValidityBitmap(int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + const int64_t actual_non_null = + valid_bits == nullptr + ? num_values + : ::arrow::internal::CountSetBits(valid_bits, valid_bits_offset, num_values); + if (ARROW_PREDICT_FALSE(actual_non_null != num_values - null_count)) { + throw ParquetException( + "Invalid validity bitmap: null_count does not match the number of non-null " + "values"); + } +} + } // namespace // ---------------------------------------------------------------------- diff --git a/cpp/src/parquet/encoding_benchmark.cc b/cpp/src/parquet/encoding_benchmark.cc index bea1a5807a2..62d8739bf03 100644 --- a/cpp/src/parquet/encoding_benchmark.cc +++ b/cpp/src/parquet/encoding_benchmark.cc @@ -1300,14 +1300,15 @@ class BenchmarkDecodeArrowByteArray : public BenchmarkDecodeArrowBasenull_bitmap_data(); total_size_ = input_array_->data()->buffers[2]->size(); @@ -1318,7 +1319,23 @@ class BenchmarkDecodeArrowByteArray : public BenchmarkDecodeArrowBase(state.range(1)) / 10000; + InitDataInputs(); + DoEncodeArrow(); + + for (auto _ : state) { + auto decoder = InitializeDecoder(); + auto acc = CreateAccumulator(); + decoder->DecodeArrow(num_values_, static_cast(input_array_->null_count()), + valid_bits_, 0, &acc); + } + state.SetBytesProcessed(state.iterations() * total_size_); + state.SetItemsProcessed(state.iterations() * num_values_); + } + protected: + double null_probability_{0.0}; std::vector values_; }; @@ -1478,6 +1495,14 @@ class BM_ArrowBinaryViewDict : public BM_ArrowBinaryDict { } }; +static void ByteArrayWithNullCustomArguments(benchmark::internal::Benchmark* b) { + b->ArgsProduct({ + benchmark::CreateRange(MIN_RANGE, MAX_RANGE, /*multi=*/4), + {1000, 5000}, + }) + ->ArgNames({"num_values", "null_in_ten_thousand"}); +} + BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, EncodeArrow) (benchmark::State& state) { EncodeArrowBenchmark(state); } BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, EncodeArrow)->Range(1 << 18, 1 << 20); @@ -1507,6 +1532,11 @@ BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrow_Dense)(benchmark::State& stat } BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrow_Dense)->Range(MIN_RANGE, MAX_RANGE); +BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrowWithNull_Dense) +(benchmark::State& state) { DecodeArrowWithNullDenseBenchmark(state); } +BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrowWithNull_Dense) + ->Apply(ByteArrayWithNullCustomArguments); + BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrowNonNull_Dense) (benchmark::State& state) { DecodeArrowNonNullDenseBenchmark(state); } BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrowNonNull_Dense) @@ -1527,6 +1557,11 @@ BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrow_Dense)(benchmark::State& BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrow_Dense) ->Range(MIN_RANGE, MAX_RANGE); +BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrowWithNull_Dense) +(benchmark::State& state) { DecodeArrowWithNullDenseBenchmark(state); } +BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrowWithNull_Dense) + ->Apply(ByteArrayWithNullCustomArguments); + BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrowNonNull_Dense) (benchmark::State& state) { DecodeArrowNonNullDenseBenchmark(state); } BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrowNonNull_Dense) diff --git a/cpp/src/parquet/encoding_test.cc b/cpp/src/parquet/encoding_test.cc index 831829e4a21..79405739408 100644 --- a/cpp/src/parquet/encoding_test.cc +++ b/cpp/src/parquet/encoding_test.cc @@ -41,6 +41,7 @@ #include "arrow/util/bitmap_writer.h" #include "arrow/util/checked_cast.h" #include "arrow/util/endian.h" +#include "arrow/util/rle_encoding_internal.h" #include "arrow/util/string.h" #include "parquet/encoding.h" #include "parquet/platform.h" @@ -1265,6 +1266,82 @@ TEST(DictEncodingAdHoc, ArrowBinaryDirectPut) { ::arrow::AssertArraysEqual(*values, *result); } +TEST(DictEncodingAdHoc, DenseDecodeRejectsInvalidInput) { + auto dictionary = ::arrow::ArrayFromJSON(::arrow::binary(), R"(["a", "bb", "ccc"])"); + auto owned_encoder = + MakeTypedEncoder(Encoding::PLAIN, /*use_dictionary=*/true); + auto* encoder = dynamic_cast*>(owned_encoder.get()); + ASSERT_NE(encoder, nullptr); + ASSERT_NO_THROW(encoder->PutDictionary(*dictionary)); + + auto dictionary_buffer = + AllocateBuffer(default_memory_pool(), encoder->dict_encoded_size()); + encoder->WriteDict(dictionary_buffer->mutable_data()); + + auto dictionary_decoder = + MakeTypedDecoder(Encoding::PLAIN, /*descr=*/nullptr); + dictionary_decoder->SetData(encoder->num_entries(), dictionary_buffer->data(), + static_cast(dictionary_buffer->size())); + + auto decoder = MakeDictDecoder(); + decoder->SetDict(dictionary_decoder.get()); + + auto ExpectDecodeFailure = [&](const uint8_t* data, int size) { + decoder->SetData(/*num_values=*/1, data, size); + typename EncodingTraits::Accumulator acc; + acc.builder = std::make_unique<::arrow::BinaryBuilder>(); + ASSERT_THROW( + decoder->DecodeArrow(/*num_values=*/1, /*null_count=*/0, + /*valid_bits=*/nullptr, /*valid_bits_offset=*/0, &acc), + ParquetException); + }; + + constexpr int kBitWidth = 2; + std::vector invalid_index_data( + 1 + ::arrow::util::RleBitPackedEncoder::MaxBufferSize(kBitWidth, /*num_values=*/1) + + ::arrow::util::RleBitPackedEncoder::MinBufferSize(kBitWidth)); + invalid_index_data[0] = kBitWidth; + ::arrow::util::RleBitPackedEncoder index_encoder( + invalid_index_data.data() + 1, static_cast(invalid_index_data.size() - 1), + kBitWidth); + ASSERT_TRUE(index_encoder.Put(/*value=*/3)); + const int invalid_index_size = 1 + index_encoder.Flush(); + ExpectDecodeFailure(invalid_index_data.data(), invalid_index_size); + + const uint8_t truncated_index_data[] = {kBitWidth}; + ExpectDecodeFailure(truncated_index_data, sizeof(truncated_index_data)); + + auto valid_indices = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 1])"); + ASSERT_NO_THROW(encoder->PutIndices(*valid_indices)); + auto valid_index_data = encoder->FlushValues(); + + auto ExpectValidityBitmapFailure = [&](const char* case_name, int num_values, + int null_count, uint8_t valid_bits) { + SCOPED_TRACE(case_name); + for (const auto& type : {::arrow::binary(), ::arrow::utf8(), ::arrow::large_binary(), + ::arrow::large_utf8()}) { + SCOPED_TRACE(type->ToString()); + decoder->SetData(num_values, valid_index_data->data(), + static_cast(valid_index_data->size())); + typename EncodingTraits::Accumulator acc; + ASSERT_OK_AND_ASSIGN(acc.builder, ::arrow::MakeBuilder(type)); + ASSERT_THROW(decoder->DecodeArrow(num_values, null_count, &valid_bits, + /*valid_bits_offset=*/0, &acc), + ParquetException); + } + }; + + // More non-null slots than promised by null_count must not read past the + // predecoded indices buffer. + ExpectValidityBitmapFailure("more non-null slots", /*num_values=*/2, + /*null_count=*/1, + /*valid_bits=*/0b00000011); + // Fewer non-null slots must not silently leave predecoded indices unconsumed. + ExpectValidityBitmapFailure("fewer non-null slots", /*num_values=*/3, + /*null_count=*/1, + /*valid_bits=*/0b00000001); +} + TEST(DictEncodingAdHoc, PutDictionaryPutIndices) { // Part of ARROW-3246 auto dict_values =