From a9ee56dcb5e4ee2f6dd42a524d95713f5c0f07d1 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 28 Jul 2026 10:52:50 +0200 Subject: [PATCH 1/2] Check level range in decoder --- cpp/src/arrow/util/rle_encoding_internal.h | 70 ++++++++++++++++------ cpp/src/parquet/column_reader.cc | 25 ++++---- 2 files changed, 65 insertions(+), 30 deletions(-) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd6530..ea4be191abc 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include "arrow/util/bit_run_reader.h" @@ -93,6 +94,9 @@ namespace arrow::util { /// It can therefore be referred to as a "typical" size for Rle and BitPacked logic. using rle_size_t = int32_t; +/// A default placeholder type. +struct RleDefault {}; + template class RleRunDecoder; @@ -322,19 +326,27 @@ class RleRunDecoder { } /// Get the next value and return false if there are no more. - [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width) { - return GetBatch(out_value, 1, value_bit_width) == 1; + template + [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width, + Func&& validator = {}) { + return GetBatch(out_value, 1, value_bit_width, std::forward(validator)) == 1; } /// Get a batch of values return the number of decoded elements. /// May write fewer elements to the output than requested if there are not enough values /// left. + template [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size, - rle_size_t value_bit_width) { + rle_size_t value_bit_width, Func&& validator = {}) { if (ARROW_PREDICT_FALSE(remaining_count_ == 0)) { return 0; } + if constexpr (!std::is_same_v, RleDefault>) { + if (ARROW_PREDICT_TRUE(batch_size > 0)) { + validator(value_); + } + } const auto to_read = std::min(remaining_count_, batch_size); std::fill(out, out + to_read, value_); remaining_count_ -= to_read; @@ -408,15 +420,18 @@ class BitPackedRunDecoder { } /// Get the next value and return false if there are no more. - [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width) { - return GetBatch(out_value, 1, value_bit_width) == 1; + template + [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width, + Func&& validator = {}) { + return GetBatch(out_value, 1, value_bit_width, std::forward(validator)) == 1; } /// Get a batch of values return the number of decoded elements. /// May write fewer elements to the output than requested if there are not enough values /// left. + template [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size, - rle_size_t value_bit_width) { + rle_size_t value_bit_width, Func&& validator = {}) { const int64_t bits_read = static_cast(values_read_) * value_bit_width; const int64_t bytes_fully_read = bits_read / 8; // The parser only creates runs whose full payload fits in max_read_bytes_ (see @@ -434,11 +449,18 @@ class BitPackedRunDecoder { if constexpr (std::is_same_v) { ::arrow::internal::unpack(unread_data, out, opts); - } else { ::arrow::internal::unpack( unread_data, reinterpret_cast*>(out), opts); } + + // Validate decoded output if given a validator + if constexpr (!std::is_same_v, RleDefault>) { + for (rle_size_t k = 0; k < opts.batch_size; ++k) { + validator(out[k]); + } + } + values_read_ += opts.batch_size; return opts.batch_size; } @@ -509,12 +531,15 @@ class RleBitPackedDecoder { /// input with zeros. Since the encoding does not differentiate between /// input values and padding, Get() returns true even for these padding /// values. - [[nodiscard]] bool Get(value_type* val); + template + [[nodiscard]] bool Get(value_type* val, Func&& validator = {}); /// Get a batch of values return the number of decoded elements. /// May write fewer elements to the output than requested if there are not enough values /// left or if an error occurred. - [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size); + template + [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size, + Func&& validator = {}); /// Like GetBatch but add spacing for null entries. /// @@ -637,11 +662,17 @@ class BitPackedDecoder : private BitPackedRunDecoder { } /// Gets the next value or returns false if there are no more or an error occurred. - [[nodiscard]] bool Get(value_type* val) { return Base::Get(val, value_bit_width_); } + template + [[nodiscard]] bool Get(value_type* val, Func&& validator = {}) { + return Base::Get(val, value_bit_width_, std::forward(validator)); + } /// Get a batch of values return the number of decoded elements. - [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size) { - return Base::GetBatch(out, batch_size, value_bit_width_); + template + [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size, + Func&& validator = {}) { + return Base::GetBatch(out, batch_size, value_bit_width_, + std::forward(validator)); } private: @@ -989,16 +1020,19 @@ RleCountUpToResult RleBitPackedDecoder::CountUpTo(value_type value, } template -bool RleBitPackedDecoder::Get(value_type* val) { - return GetBatch(val, 1) == 1; +template +bool RleBitPackedDecoder::Get(value_type* val, Func&& validator) { + return GetBatch(val, 1, std::forward(validator)) == 1; } template -auto RleBitPackedDecoder::GetBatch(value_type* out, - rle_size_t batch_size) -> rle_size_t { +template +auto RleBitPackedDecoder::GetBatch(value_type* out, rle_size_t batch_size, + Func&& validator) -> rle_size_t { return ProcessValues( - [&out, this](auto& decoder, rle_size_t run_batch_size) { - const auto read = decoder.GetBatch(out, run_batch_size, value_bit_width_); + [&out, this, &validator](auto& decoder, rle_size_t run_batch_size) { + const auto read = + decoder.GetBatch(out, run_batch_size, value_bit_width_, validator); out += read; return read; }, diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 162a72bd157..4adda171228 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -106,8 +106,10 @@ struct LevelDecoder::Impl { std::variant decoder = {}; - [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { - return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); + template + [[nodiscard]] int GetBatch(int16_t* out, int batch_size, Func&& validator) { + return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size, validator); }, + decoder); } [[nodiscard]] int Advance(int batch_size) { @@ -189,16 +191,15 @@ void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, int LevelDecoder::Decode(int batch_size, int16_t* levels) { const int num_values = std::min(num_values_remaining_, batch_size); - const int num_decoded = impl_->GetBatch(levels, num_values); - if (num_decoded > 0) { - internal::MinMax min_max = internal::FindMinMax(levels, num_decoded); - if (ARROW_PREDICT_FALSE(min_max.min < 0 || min_max.max > max_level_)) { - std::stringstream ss; - ss << "Malformed levels. min: " << min_max.min << " max: " << min_max.max - << " out of range. Max Level: " << max_level_; - throw ParquetException(ss.str()); - } - } + const int num_decoded = + impl_->GetBatch(levels, num_values, [max = max_level_](int16_t value) { + if (ARROW_PREDICT_FALSE(value < 0 || value > max)) { + std::stringstream ss; + ss << "Malformed level with value " << value << " is not in the range [0, " + << max << "]"; + throw ParquetException(ss.str()); + } + }); num_values_remaining_ -= num_decoded; return num_decoded; } From 98f1769cebb04c68656b6bcc01c33d77bc4588f9 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 29 Jul 2026 09:20:19 +0200 Subject: [PATCH 2/2] Use FindMinMax --- cpp/src/arrow/util/rle_encoding_internal.h | 6 ++---- cpp/src/parquet/column_reader.cc | 11 ++++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index ea4be191abc..69fdb00cc1b 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -344,7 +344,7 @@ class RleRunDecoder { if constexpr (!std::is_same_v, RleDefault>) { if (ARROW_PREDICT_TRUE(batch_size > 0)) { - validator(value_); + validator(static_cast(&value_), 1); } } const auto to_read = std::min(remaining_count_, batch_size); @@ -456,9 +456,7 @@ class BitPackedRunDecoder { // Validate decoded output if given a validator if constexpr (!std::is_same_v, RleDefault>) { - for (rle_size_t k = 0; k < opts.batch_size; ++k) { - validator(out[k]); - } + validator(static_cast(out), opts.batch_size); } values_read_ += opts.batch_size; diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 4adda171228..9889ed308d0 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -191,12 +191,13 @@ void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, int LevelDecoder::Decode(int batch_size, int16_t* levels) { const int num_values = std::min(num_values_remaining_, batch_size); - const int num_decoded = - impl_->GetBatch(levels, num_values, [max = max_level_](int16_t value) { - if (ARROW_PREDICT_FALSE(value < 0 || value > max)) { + const int num_decoded = impl_->GetBatch( + levels, num_values, [max = max_level_](const auto* levels, auto size) { + internal::MinMax min_max = internal::FindMinMax(levels, size); + if (ARROW_PREDICT_FALSE(min_max.min < 0 || min_max.max > max)) { std::stringstream ss; - ss << "Malformed level with value " << value << " is not in the range [0, " - << max << "]"; + ss << "Malformed levels. min: " << min_max.min << " max: " << min_max.max + << " out of range. Max Level: " << max; throw ParquetException(ss.str()); } });