From c38f9c61aaf25195dd92ac4f5fddb02324d97ff9 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 24 Jun 2026 11:22:26 +0200 Subject: [PATCH 01/13] Simplify LevelReader --- cpp/src/parquet/column_reader.cc | 127 ++++++++++++++++++++++--------- cpp/src/parquet/column_reader.h | 33 +++----- 2 files changed, 103 insertions(+), 57 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 7241632e4b62..c21288134c0d 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "arrow/array.h" @@ -95,7 +96,77 @@ constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = } // namespace -LevelDecoder::LevelDecoder() : num_values_remaining_(0) {} +/****************** + * LevelDecoder * + ******************/ + +namespace { + +/// Adapter around ``BitReader`` that mimics the API of ``RleBitPackedDecoder``. +/// +/// Best would be to make this a first class citizen, possibly reusing +/// ``BitPackedRunDecoder`` but the preconditions on the possible number of values +/// that the run can represent differ. +template +class BitPackedDecoderWrapper : private ::arrow::bit_util::BitReader { + public: + /// The type in which the data should be decoded. + using value_type = T; + using rle_size_t = ::arrow::util::rle_size_t; + + BitPackedDecoderWrapper() noexcept = default; + + BitPackedDecoderWrapper(const uint8_t* data, rle_size_t data_size, + rle_size_t value_bit_width) noexcept { + Reset(data, data_size, value_bit_width); + } + + void Reset(const uint8_t* data, rle_size_t data_size, + rle_size_t value_bit_width) noexcept { + value_bit_width_ = value_bit_width; + return Base::Reset(data, data_size); + } + + /// Whether there is still values to iterate over. + bool exhausted() const { return Base::bytes_left() >= value_bit_width_; } + + /// Gets the next value or returns false if there are no more or an error occurred. + [[nodiscard]] bool Get(value_type* val) { + return Base::GetValue(value_bit_width_, val); + } + + /// 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(value_bit_width_, out, batch_size); + } + + private: + using Base = ::arrow::bit_util::BitReader; + + rle_size_t value_bit_width_ = {}; +}; +} // namespace + +struct LevelDecoder::Impl { + using RleBitPackedDecoder = ::arrow::util::RleBitPackedDecoder; + using BitPackedDecoder = BitPackedDecoderWrapper; + + std::variant decoder = {}; + + constexpr bool is_rle_bit_packed() const { + return std::holds_alternative(decoder); + } + + constexpr bool is_bit_packed() const { + return std::holds_alternative(decoder); + } + + [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { + return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); + } +}; + +LevelDecoder::LevelDecoder() : decoder_(std::make_unique()) {} LevelDecoder::~LevelDecoder() = default; @@ -103,44 +174,38 @@ int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, int num_buffered_values, const uint8_t* data, int32_t data_size) { max_level_ = max_level; - int32_t num_bytes = 0; - encoding_ = encoding; num_values_remaining_ = num_buffered_values; - bit_width_ = bit_util::Log2(max_level + 1); + const int value_bit_width = bit_util::Log2(max_level + 1); + switch (encoding) { case Encoding::RLE: { if (data_size < 4) { throw ParquetException("Received invalid levels (corrupt data page?)"); } - num_bytes = ::arrow::util::SafeLoadAs(data); + const auto num_bytes = ::arrow::util::SafeLoadAs(data); if (num_bytes < 0 || num_bytes > data_size - 4) { throw ParquetException("Received invalid number of bytes (corrupt data page?)"); } - const uint8_t* decoder_data = data + 4; - if (!rle_decoder_) { - rle_decoder_ = std::make_unique<::arrow::util::RleBitPackedDecoder>( - decoder_data, num_bytes, bit_width_); - } else { - rle_decoder_->Reset(decoder_data, num_bytes, bit_width_); - } + this->decoder_->decoder = Impl::RleBitPackedDecoder( // + /* data= */ data + 4, + /* data_size =*/num_bytes, + /* value_bit_width= */ value_bit_width); return 4 + num_bytes; } case Encoding::BIT_PACKED: { int num_bits = 0; - if (MultiplyWithOverflow(num_buffered_values, bit_width_, &num_bits)) { + if (MultiplyWithOverflow(num_buffered_values, value_bit_width, &num_bits)) { throw ParquetException( "Number of buffered values too large (corrupt data page?)"); } - num_bytes = static_cast(bit_util::BytesForBits(num_bits)); + const auto num_bytes = static_cast(bit_util::BytesForBits(num_bits)); if (num_bytes < 0 || num_bytes > data_size) { throw ParquetException("Received invalid number of bytes (corrupt data page?)"); } - if (!bit_packed_decoder_) { - bit_packed_decoder_ = - std::make_unique<::arrow::bit_util::BitReader>(data, num_bytes); - } else { - bit_packed_decoder_->Reset(data, num_bytes); - } + this->decoder_->decoder = Impl::BitPackedDecoder( // + /* data= */ data, + /* data_size =*/num_bytes, + /* value_bit_width= */ value_bit_width); return num_bytes; } default: @@ -157,27 +222,17 @@ void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, if (num_bytes < 0) { throw ParquetException("Invalid page header (corrupt data page?)"); } - encoding_ = Encoding::RLE; num_values_remaining_ = num_buffered_values; - bit_width_ = bit_util::Log2(max_level + 1); - if (!rle_decoder_) { - rle_decoder_ = std::make_unique<::arrow::util::RleBitPackedDecoder>( - data, num_bytes, bit_width_); - } else { - rle_decoder_->Reset(data, num_bytes, bit_width_); - } + this->decoder_->decoder = Impl::RleBitPackedDecoder( // + /* data= */ data, + /* data_size =*/num_bytes, + /* value_bit_width= */ bit_util::Log2(max_level + 1)); } int LevelDecoder::Decode(int batch_size, int16_t* levels) { - int num_decoded = 0; - - int num_values = std::min(num_values_remaining_, batch_size); - if (encoding_ == Encoding::RLE) { - num_decoded = rle_decoder_->GetBatch(levels, num_values); - } else { - num_decoded = bit_packed_decoder_->GetBatch(bit_width_, levels, num_values); - } + const int num_values = std::min(num_values_remaining_, batch_size); + const int num_decoded = decoder_->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_)) { diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index ad20d4a9e6e6..b8a869c42c9b 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -32,19 +32,6 @@ #include "parquet/schema.h" #include "parquet/types.h" -namespace arrow { - -namespace bit_util { -class BitReader; -} // namespace bit_util - -namespace util { -template -class RleBitPackedDecoder; -} // namespace util - -} // namespace arrow - namespace parquet { class Decryptor; @@ -79,25 +66,29 @@ struct PARQUET_EXPORT DataPageStats { class PARQUET_EXPORT LevelDecoder { public: LevelDecoder(); + ~LevelDecoder(); - // Initialize the LevelDecoder state with new data - // and return the number of bytes consumed + /// Initialize the LevelDecoder state with new data from a legacy (V1) page. + /// + /// @return the number of bytes consumed int SetData(Encoding::type encoding, int16_t max_level, int num_buffered_values, const uint8_t* data, int32_t data_size); + /// Initialize the LevelDecoder state with new data from a V2 page. + /// + /// Encoding of the level in the void SetDataV2(int32_t num_bytes, int16_t max_level, int num_buffered_values, const uint8_t* data); - // Decodes a batch of levels into an array and returns the number of levels decoded + /// Decode a batch of levels into an array and returns the number of levels decoded int Decode(int batch_size, int16_t* levels); private: - int bit_width_; - int num_values_remaining_; - Encoding::type encoding_; - std::unique_ptr<::arrow::util::RleBitPackedDecoder> rle_decoder_; - std::unique_ptr<::arrow::bit_util::BitReader> bit_packed_decoder_; + struct Impl; + + std::unique_ptr decoder_; + int num_values_remaining_ = 0; int16_t max_level_; }; From a5104acc696675ad9bf5ba8a2335558c130fbd86 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 24 Jun 2026 15:55:59 +0200 Subject: [PATCH 02/13] Proper TypedColumn def/rep level accessors --- cpp/src/parquet/column_reader.cc | 83 +++++++++++++++----------------- 1 file changed, 39 insertions(+), 44 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index c21288134c0d..892e1a5b9ffc 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -693,14 +693,7 @@ class ColumnReaderImplBase { using T = typename DType::c_type; ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - max_def_level_(descr->max_definition_level()), - max_rep_level_(descr->max_repetition_level()), - num_buffered_values_(0), - num_decoded_values_(0), - pool_(pool), - current_decoder_(nullptr), - current_encoding_(Encoding::UNKNOWN) {} + : descr_(descr), pool_(pool) {} virtual ~ColumnReaderImplBase() = default; @@ -730,7 +723,7 @@ class ColumnReaderImplBase { // // Returns the number of decoded definition levels int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels) { - if (max_def_level_ == 0) { + if (max_def_level() == 0) { return 0; } return definition_level_decoder_.Decode(static_cast(batch_size), levels); @@ -750,7 +743,7 @@ class ColumnReaderImplBase { // Read multiple repetition levels into preallocated memory // Returns the number of decoded repetition levels int64_t ReadRepetitionLevels(int64_t batch_size, int16_t* levels) { - if (max_rep_level_ == 0) { + if (max_rep_level() == 0) { return 0; } return repetition_level_decoder_.Decode(static_cast(batch_size), levels); @@ -846,9 +839,9 @@ class ColumnReaderImplBase { // Data page Layout: Repetition Levels - Definition Levels - encoded values. // Levels are encoded as rle or bit-packed. // Init repetition levels - if (max_rep_level_ > 0) { + if (max_rep_level() > 0) { int32_t rep_levels_bytes = repetition_level_decoder_.SetData( - repetition_level_encoding, max_rep_level_, + repetition_level_encoding, max_rep_level(), static_cast(num_buffered_values_), buffer, max_size); buffer += rep_levels_bytes; levels_byte_size += rep_levels_bytes; @@ -858,9 +851,9 @@ class ColumnReaderImplBase { // if the initial value is invalid // Init definition levels - if (max_def_level_ > 0) { + if (max_def_level() > 0) { int32_t def_levels_bytes = definition_level_decoder_.SetData( - definition_level_encoding, max_def_level_, + definition_level_encoding, max_def_level(), static_cast(num_buffered_values_), buffer, max_size); levels_byte_size += def_levels_bytes; max_size -= def_levels_bytes; @@ -885,9 +878,9 @@ class ColumnReaderImplBase { throw ParquetException("Data page too small for levels (corrupt header?)"); } - if (max_rep_level_ > 0) { + if (max_rep_level() > 0) { repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(), - max_rep_level_, + max_rep_level(), static_cast(num_buffered_values_), buffer); } // ARROW-17453: Even if max_rep_level_ is 0, there may still be @@ -895,9 +888,9 @@ class ColumnReaderImplBase { // some writers (e.g. Athena) buffer += page.repetition_levels_byte_length(); - if (max_def_level_ > 0) { + if (max_def_level() > 0) { definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(), - max_def_level_, + max_def_level(), static_cast(num_buffered_values_), buffer); } @@ -957,9 +950,11 @@ class ColumnReaderImplBase { return num_buffered_values_ - num_decoded_values_; } + int16_t max_def_level() const { return descr_->max_definition_level(); } + + int16_t max_rep_level() const { return descr_->max_repetition_level(); } + const ColumnDescriptor* descr_; - const int16_t max_def_level_; - const int16_t max_rep_level_; std::unique_ptr pager_; std::shared_ptr current_page_; @@ -976,17 +971,17 @@ class ColumnReaderImplBase { // values. For repeated or optional values, there may be fewer data values // than levels, and this tells you how many encoded levels there are in that // case. - int64_t num_buffered_values_; + int64_t num_buffered_values_ = 0; // The number of values from the current data page that have been decoded // into memory or skipped over. - int64_t num_decoded_values_; + int64_t num_decoded_values_ = 0; ::arrow::MemoryPool* pool_; using DecoderType = TypedDecoder; - DecoderType* current_decoder_; - Encoding::type current_encoding_; + DecoderType* current_decoder_ = nullptr; + Encoding::type current_encoding_ = Encoding::UNKNOWN; /// Flag to signal when a new dictionary has been set, for the benefit of /// DictionaryRecordReader @@ -1075,7 +1070,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, batch_size = std::min(batch_size, this->available_values_current_page()); // If the field is required and non-repeated, there are no definition levels - if (this->max_def_level_ > 0 && def_levels != nullptr) { + if (this->max_def_level() > 0 && def_levels != nullptr) { *num_def_levels = this->ReadDefinitionLevels(batch_size, def_levels); if (ARROW_PREDICT_FALSE(*num_def_levels != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1083,7 +1078,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, // TODO(wesm): this tallying of values-to-decode can be performed with better // cache-efficiency if fused with the level decoding. *non_null_values_to_read += - std::count(def_levels, def_levels + *num_def_levels, this->max_def_level_); + std::count(def_levels, def_levels + *num_def_levels, this->max_def_level()); } else { // Required field, read all values if (num_def_levels != nullptr) { @@ -1093,7 +1088,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, } // Not present for non-repeated fields - if (this->max_rep_level_ > 0 && rep_levels != nullptr) { + if (this->max_rep_level() > 0 && rep_levels != nullptr) { int64_t num_rep_levels = this->ReadRepetitionLevels(batch_size, rep_levels); if (batch_size != num_rep_levels) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1378,7 +1373,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, break; } - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { ReserveLevels(batch_size); int16_t* def_levels = this->def_levels() + levels_written_; @@ -1388,7 +1383,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { int64_t rep_levels_read = this->ReadRepetitionLevels(batch_size, rep_levels); if (ARROW_PREDICT_FALSE(rep_levels_read != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1414,7 +1409,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, void ThrowAwayLevels(int64_t start_levels_position) { ARROW_DCHECK_LE(levels_position_, levels_written_); ARROW_DCHECK_LE(start_levels_position, levels_position_); - ARROW_DCHECK_GT(this->max_def_level_, 0); + ARROW_DCHECK_GT(this->max_def_level(), 0); ARROW_DCHECK_NE(def_levels_, nullptr); int64_t gap = levels_position_ - start_levels_position; @@ -1432,7 +1427,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, left_shift(def_levels_.get()); - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { ARROW_DCHECK_NE(rep_levels_, nullptr); left_shift(rep_levels_.get()); } @@ -1445,7 +1440,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Skip records that we have in our buffer. This function is only for // non-repeated fields. int64_t SkipRecordsInBufferNonRepeated(int64_t num_records) { - ARROW_DCHECK_EQ(this->max_rep_level_, 0); + ARROW_DCHECK_EQ(this->max_rep_level(), 0); if (!this->has_values_to_process() || num_records == 0) return 0; int64_t remaining_records = levels_written_ - levels_position_; @@ -1459,7 +1454,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // First we need to figure out how many present/not-null values there are. int64_t values_to_read = std::count(def_levels() + start_levels_position, def_levels() + levels_position_, - this->max_def_level_); + this->max_def_level()); // Now that we have figured out number of values to read, we do not need // these levels anymore. We will remove these values from the buffer. @@ -1506,7 +1501,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // desired number of records or we run out of values in the column chunk. // Returns number of skipped records. int64_t SkipRecordsRepeated(int64_t num_records) { - ARROW_DCHECK_GT(this->max_rep_level_, 0); + ARROW_DCHECK_GT(this->max_rep_level(), 0); int64_t skipped_records = 0; // First consume what is in the buffer. @@ -1596,11 +1591,11 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Top level required field. Number of records equals to number of levels, // and there is not read-ahead for levels. - if (this->max_rep_level_ == 0 && this->max_def_level_ == 0) { + if (this->max_rep_level() == 0 && this->max_def_level() == 0) { return this->Skip(num_records); } int64_t skipped_records = 0; - if (this->max_rep_level_ == 0) { + if (this->max_rep_level() == 0) { // Non-repeated optional field. // First consume whatever is in the buffer. skipped_records = SkipRecordsInBufferNonRepeated(num_records); @@ -1661,7 +1656,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, int64_t records_read = 0; const int16_t* const rep_levels = this->rep_levels(); const int16_t* const def_levels = this->def_levels(); - ARROW_DCHECK_GT(this->max_rep_level_, 0); + ARROW_DCHECK_GT(this->max_rep_level(), 0); // If at_record_start_ is true, we are seeing the start of a record // for the second time, such as after repeated calls to // DelimitRecords. In this case we must continue until we find @@ -1707,7 +1702,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } // Scan definition levels to find number of physical values *values_seen = std::count(def_levels + level, def_levels + levels_position_, - this->max_def_level_); + this->max_def_level()); return records_read; } @@ -1734,7 +1729,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } void ReserveLevels(int64_t extra_levels) { - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { const int64_t new_levels_capacity = UpdateCapacity(levels_capacity_, levels_written_, extra_levels); if (new_levels_capacity > levels_capacity_) { @@ -1745,7 +1740,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } PARQUET_THROW_NOT_OK( def_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { PARQUET_THROW_NOT_OK( rep_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); } @@ -1885,7 +1880,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // When reading dense we need to figure out number of values to read. const int16_t* def_levels = this->def_levels(); *values_to_read += std::count(def_levels + start_levels_position, - def_levels + levels_position_, this->max_def_level_); + def_levels + levels_position_, this->max_def_level()); ReadValuesDense(*values_to_read); } @@ -1924,11 +1919,11 @@ class TypedRecordReader : public TypedColumnReaderImpl, int64_t records_read = 0; int64_t values_to_read = 0; int64_t null_count = 0; - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { // Repeated fields may be nullable or not. // This call updates levels_position_. records_read = ReadRepeatedRecords(num_records, &values_to_read, &null_count); - } else if (this->max_def_level_ > 0) { + } else if (this->max_def_level() > 0) { // Non-repeated optional values are always nullable. // This call updates levels_position_. ARROW_DCHECK(nullable_values()); @@ -1951,7 +1946,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, null_count_ += null_count; } // Total values, including null spaces, if any - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { // Optional, repeated, or some mix thereof this->ConsumeBufferedValues(levels_position_ - start_levels_position); } else { From 9902d19564c8c0269722f3e5063905527fb11fc2 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 24 Jun 2026 16:15:23 +0200 Subject: [PATCH 03/13] in object level storage --- cpp/src/parquet/column_reader.cc | 22 ++++++++++++++++------ cpp/src/parquet/column_reader.h | 5 ++++- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 892e1a5b9ffc..79753b0339f5 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -166,7 +166,8 @@ struct LevelDecoder::Impl { } }; -LevelDecoder::LevelDecoder() : decoder_(std::make_unique()) {} +LevelDecoder::LevelDecoder(int16_t max_level) + : decoder_(std::make_unique()), max_level_(max_level) {} LevelDecoder::~LevelDecoder() = default; @@ -693,7 +694,10 @@ class ColumnReaderImplBase { using T = typename DType::c_type; ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), pool_(pool) {} + : descr_(descr), + definition_level_decoder_(descr->max_definition_level()), + repetition_level_decoder_(descr->max_repetition_level()), + pool_(pool) {} virtual ~ColumnReaderImplBase() = default; @@ -950,19 +954,25 @@ class ColumnReaderImplBase { return num_buffered_values_ - num_decoded_values_; } - int16_t max_def_level() const { return descr_->max_definition_level(); } + int16_t max_def_level() const { + // max level indirectly part of this object storage + return definition_level_decoder_.max_level(); + } - int16_t max_rep_level() const { return descr_->max_repetition_level(); } + int16_t max_rep_level() const { + // max level indirectly part of this object storage + return repetition_level_decoder_.max_level(); + } const ColumnDescriptor* descr_; std::unique_ptr pager_; std::shared_ptr current_page_; - // Not set if full schema for this field has no optional or repeated elements + // No data set if full schema for this field has no optional or repeated elements LevelDecoder definition_level_decoder_; - // Not set for flat schemas. + // No data set for flat schemas. LevelDecoder repetition_level_decoder_; // The total number of values stored in the data page. This is the maximum of diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index b8a869c42c9b..802e0062cc2f 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -65,7 +65,7 @@ struct PARQUET_EXPORT DataPageStats { class PARQUET_EXPORT LevelDecoder { public: - LevelDecoder(); + explicit LevelDecoder(int16_t max_level = 0); ~LevelDecoder(); @@ -84,6 +84,9 @@ class PARQUET_EXPORT LevelDecoder { /// Decode a batch of levels into an array and returns the number of levels decoded int Decode(int batch_size, int16_t* levels); + /// Return the max level used in this decoder. + int max_level() const { return max_level_; } + private: struct Impl; From 749d345ebafed3216abcbae21b8c83001a5b5b74 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 7 Jul 2026 15:05:33 +0200 Subject: [PATCH 04/13] Add SkippableTypedDecoder --- cpp/src/parquet/column_reader.cc | 58 ++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 79753b0339f5..ad5a9ee74968 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -685,6 +685,64 @@ std::unique_ptr PageReader::Open(std::shared_ptr s namespace { +/// Wrapper around a `TypedDecoder` pointer to skip values. +/// +/// Use a scratch buffer to decode values into that buffer. +/// This was migrated here from a historical implementation. +/// Ideally all decoders would implement a `Skip` functionality that would at best +/// avoid decoding, and at worst, decode without intermediarry allocation. +template +class SkippableTypedDecoder { + public: + using Decoder = TypedDecoder; + using T = typename Decoder::T; + + static constexpr int64_t kValueByteSize = type_traits::value_byte_size; + static constexpr int64_t kScratchValueCount = kScratchValueCount_; + static constexpr int64_t kScratchByteSize = kScratchValueCount * kValueByteSize; + + SkippableTypedDecoder() = default; + + explicit SkippableTypedDecoder(Decoder* decoder) : decoder_(decoder) {} + + void SetDecoder(Decoder* decoder) { decoder_ = decoder; } + + const Decoder* get() const { return decoder_; } + + Decoder* get() { return decoder_; } + + int64_t Skip(int64_t num_values, ::arrow::MemoryPool* pool = nullptr) { + EnsureScratch(pool); + + int64_t total_read = 0; + int iter_read = 0; + do { + static_assert(kScratchValueCount <= std::numeric_limits::max()); + const int batch_size = + static_cast(std::min(kScratchValueCount, num_values - total_read)); + + iter_read = get()->Decode(scratch_->mutable_data_as(), batch_size); + total_read += iter_read; + } while (iter_read > 0 && total_read < num_values); + + return total_read; + } + + private: + /// Scratch space to skip decode values that need skipping. + /// We actually do not need the whole shared_ptr machinery but it was historically + /// chosen for ease of use with ``AllocateBuffer`` and migrated here. + std::shared_ptr scratch_ = nullptr; + Decoder* decoder_ = nullptr; + + void EnsureScratch(::arrow::MemoryPool* pool) { + if (this->scratch_ == nullptr) { + this->scratch_ = AllocateBuffer(pool, kScratchByteSize); + } + ARROW_DCHECK_NE(this->scratch_, nullptr); + } +}; + // ---------------------------------------------------------------------- // Impl base class for TypedColumnReader and RecordReader From cfd79979b2a844b5e000c4ce73c9d6bb647ac925 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 7 Jul 2026 17:40:18 +0200 Subject: [PATCH 05/13] Add RleBitPackedDecoder::Advance --- cpp/src/arrow/util/rle_encoding_internal.h | 146 ++++++++++----------- cpp/src/arrow/util/rle_encoding_test.cc | 117 +++++++++++++---- 2 files changed, 161 insertions(+), 102 deletions(-) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 5dc94f368d3c..11f36c498e23 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -448,6 +448,10 @@ class RleBitPackedDecoder { /// This is how one can check for errors. bool exhausted() const { return (run_remaining() == 0) && parser_.exhausted(); } + /// Advance by as many values as provided or until exhaustion of the decoder. + /// Return the number of values skipped. + [[nodiscard]] rle_size_t Advance(rle_size_t batch_size); + /// Gets the next value or returns false if there are no more or an error occurred. /// /// NB: Because the encoding only supports literal runs with lengths @@ -500,12 +504,15 @@ class RleBitPackedDecoder { return std::visit([](const auto& dec) { return dec.remaining(); }, decoder_); } - /// Get a batch of values from the current run and return the number elements read. - [[nodiscard]] rle_size_t RunGetBatch(value_type* out, rle_size_t batch_size) { - return std::visit( - [&](auto& dec) { return dec.GetBatch(out, batch_size, value_bit_width_); }, - decoder_); - } + /// Process data in the current run and subsequent ones. + /// + /// `func` is called as `func(decoder, run_batch_size)` where `decoder` are + /// statically-typed run decoder (not the variant). + /// Must return the number of values it processed. + /// + /// Return the number of values processed. + template + rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size); /// Utility methods for retrieving spaced values. template @@ -782,30 +789,26 @@ struct RleBitPackedDecoderGetRunDecoder { }; template -bool RleBitPackedDecoder::Get(value_type* val) { - return GetBatch(val, 1) == 1; -} - -template -auto RleBitPackedDecoder::GetBatch(value_type* out, - rle_size_t batch_size) -> rle_size_t { +template +auto RleBitPackedDecoder::ProcessValues(Callable&& func, + rle_size_t batch_size) -> rle_size_t { using ControlFlow = RleBitPackedParser::ControlFlow; if (ARROW_PREDICT_FALSE(batch_size == 0 || exhausted())) { return 0; } - rle_size_t values_read = 0; + rle_size_t values_processed = 0; // Remaining from a previous call that would have left some unread data from a run. if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { - const auto read = RunGetBatch(out, batch_size); - values_read += read; - out += read; + const auto processed = + std::visit([&](auto& decoder) { return func(decoder, batch_size); }, decoder_); + values_processed += processed; // Either we fulfilled all the batch to be read or we finished remaining run. - if (ARROW_PREDICT_FALSE(values_read == batch_size)) { - return values_read; + if (ARROW_PREDICT_FALSE(values_processed == batch_size)) { + return values_processed; } ARROW_DCHECK(run_remaining() == 0); } @@ -814,23 +817,47 @@ auto RleBitPackedDecoder::GetBatch(value_type* out, using RunDecoder = typename RleBitPackedDecoderGetRunDecoder::type; - ARROW_DCHECK_LT(values_read, batch_size); + ARROW_DCHECK_LT(values_processed, batch_size); + // Local decoder to keep its type transparent to the compiler. RunDecoder decoder(run, value_bit_width_); - const auto read = decoder.GetBatch(out, batch_size - values_read, value_bit_width_); - ARROW_DCHECK_LE(read, batch_size - values_read); - values_read += read; - out += read; + const auto read = func(decoder, batch_size - values_processed); + ARROW_DCHECK_LE(read, batch_size - values_processed); + values_processed += read; - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(values_read == batch_size || read == 0)) { + if (ARROW_PREDICT_FALSE(values_processed == batch_size || read == 0)) { decoder_ = std::move(decoder); return ControlFlow::Break; } - return ControlFlow::Continue; }); - return values_read; + return values_processed; +} + +template +auto RleBitPackedDecoder::Advance(rle_size_t batch_size) -> rle_size_t { + return ProcessValues( + [](auto& decoder, rle_size_t run_batch_size) { + return decoder.Advance(run_batch_size); + }, + batch_size); +} + +template +bool RleBitPackedDecoder::Get(value_type* val) { + return GetBatch(val, 1) == 1; +} + +template +auto RleBitPackedDecoder::GetBatch(value_type* out, + rle_size_t batch_size) -> 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 += read; + return read; + }, + batch_size); } namespace internal { @@ -1269,8 +1296,6 @@ template auto RleBitPackedDecoder::GetBatchWithDict(const V* dictionary, int32_t dictionary_length, V* out, rle_size_t batch_size) -> rle_size_t { - using ControlFlow = RleBitPackedParser::ControlFlow; - if (ARROW_PREDICT_FALSE(batch_size <= 0 || dictionary_length == 0)) { // Either empty batch or invalid dictionary return 0; @@ -1278,59 +1303,20 @@ auto RleBitPackedDecoder::GetBatchWithDict(const V* dictionary, internal::DictionaryConverter converter{dictionary, dictionary_length}; - // Make lightweight BitRun class to reuse previous methods. + // Make lightweight BitRun class to reuse the spaced code path with no nulls. constexpr internal::UnreachableBitRunReader validity_reader{}; internal::AllSetBitRun validity_run = {batch_size}; - rle_size_t values_read = 0; - auto batch_values_remaining = [&]() { - ARROW_DCHECK_LE(values_read, batch_size); - return batch_size - values_read; - }; - - if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { - const auto read = internal::RunGetSpaced(&converter, out, batch_size, - /* null_count= */ 0, value_bit_width_, - &validity_reader, &validity_run, &decoder_); - - ARROW_DCHECK_EQ(read.null_read, 0); - values_read += read.values_read; - out += read.values_read; - - // Either we fulfilled all the batch values to be read - if (ARROW_PREDICT_FALSE(values_read >= batch_size)) { - // There may be remaining null if they are not greedily filled - return values_read; - } - - // We finished the remaining run - ARROW_DCHECK(run_remaining() == 0); - } - - parser_.ParseWithCallable([&](auto run) { - using RunDecoder = - typename RleBitPackedDecoderGetRunDecoder::type; - - RunDecoder decoder(run, value_bit_width_); - - const auto read = internal::RunGetSpaced(&converter, out, batch_values_remaining(), - /* null_count= */ 0, value_bit_width_, - &validity_reader, &validity_run, &decoder); - - ARROW_DCHECK_EQ(read.null_read, 0); - values_read += read.values_read; - out += read.values_read; - - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(read.values_read == 0 || values_read == batch_size)) { - decoder_ = std::move(decoder); - return ControlFlow::Break; - } - - return ControlFlow::Continue; - }); - - return values_read; + return ProcessValues( + [&](auto& decoder, rle_size_t run_batch_size) { + const auto read = internal::RunGetSpaced( + &converter, out, run_batch_size, /* null_count= */ 0, value_bit_width_, + &validity_reader, &validity_run, &decoder); + ARROW_DCHECK_EQ(read.null_read, 0); + out += read.values_read; + return read.values_read; + }, + batch_size); } template diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index 788450c67c90..a371f7514441 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -989,6 +989,39 @@ TEST(BitRle, Overflow) { } } +/// Encode the values of an Array with RleBitPackedEncoder. +/// +/// @param spaced If false, treat Nulls as regular data (i.e. their raw value is +/// encoded). If true, Nulls are skipped and only valid values are encoded. +template +std::vector EncodeTestArray(const Array& data, int bit_width, + bool spaced = false) { + using ArrayType = typename TypeTraits::ArrayType; + + const auto data_size = static_cast(data.length()); + const auto values_count = + static_cast(data.length() - (spaced ? data.null_count() : 0)); + const int buffer_size = + static_cast(RleBitPackedEncoder::MaxBufferSize(bit_width, values_count) + + RleBitPackedEncoder::MinBufferSize(bit_width)); + const auto* data_values = static_cast(data).raw_values(); + + std::vector buffer(buffer_size); + RleBitPackedEncoder encoder(buffer.data(), buffer_size, bit_width); + rle_size_t encoded_values_count = 0; + for (rle_size_t i = 0; i < data_size; ++i) { + if (data.IsValid(i) || !spaced) { + EXPECT_TRUE(encoder.Put(static_cast(data_values[i]))) + << "Encoding failed in pos " << i << ", current encoder len: " << encoder.len(); + ++encoded_values_count; + } + } + EXPECT_EQ(encoded_values_count, values_count) + << "All values input were not encoded successfully by the encoder"; + buffer.resize(encoder.Flush()); + return buffer; +} + /// Check RleBitPacked encoding/decoding round trip. /// /// \param spaced If set to false, treat Nulls in the input array as regular data. @@ -1002,38 +1035,20 @@ void CheckRoundTrip(const Array& data, int bit_width, bool spaced, int32_t parts using value_type = typename Type::c_type; const int data_size = static_cast(data.length()); - const int data_values_count = - static_cast(data.length() - spaced * data.null_count()); - const int buffer_size = static_cast( - ::arrow::util::RleBitPackedEncoder::MaxBufferSize(bit_width, data_values_count) + - ::arrow::util::RleBitPackedEncoder::MinBufferSize(bit_width)); ASSERT_GE(parts, 1); ASSERT_LE(parts, data_size); ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", spaced = ", spaced, - ", data_size = ", data_size, ", buffer_size = ", buffer_size); + ", data_size = ", data_size); const value_type* data_values = static_cast(data).raw_values(); // Encode the data into `buffer` using the encoder. - std::vector buffer(buffer_size); - RleBitPackedEncoder encoder(buffer.data(), buffer_size, bit_width); - int32_t encoded_values_size = 0; - for (int i = 0; i < data_size; ++i) { - // Depending on `spaced` we treat nulls as regular values. - if (data.IsValid(i) || !spaced) { - bool success = encoder.Put(static_cast(data_values[i])); - ASSERT_TRUE(success) << "Encoding failed in pos " << i - << ", current encoder len: " << encoder.len(); - ++encoded_values_size; - } - } - int encoded_byte_size = encoder.Flush(); - ASSERT_EQ(encoded_values_size, data_values_count) - << "All values input were not encoded successfully by the encoder"; + const auto buffer = EncodeTestArray(data, bit_width, spaced); // Now we verify batch read - RleBitPackedDecoder decoder(buffer.data(), encoded_byte_size, bit_width); + RleBitPackedDecoder decoder(buffer.data(), static_cast(buffer.size()), + bit_width); // We will only use one of them depending on whether this is a dictionary tests std::vector dict_read; std::vector values_read; @@ -1102,6 +1117,60 @@ void CheckRoundTrip(const Array& data, int bit_width, bool spaced, int32_t parts } } +/// Check RleBitPackedDecoder::Advance, which spans multiple runs through the +/// parser (unlike the per-run decoder Advance). +/// +/// Nulls are treated as regular data (i.e. their raw value is encoded and +/// decoded). Advance and GetBatch are interleaved so that run boundaries are +/// crossed and partial runs consumed, verifying decoded values as we go. +template +void CheckAdvance(const Array& data, int bit_width) { + using ArrayType = typename TypeTraits::ArrayType; + using value_type = typename Type::c_type; + + const auto data_size = static_cast(data.length()); + const value_type* data_values = static_cast(data).raw_values(); + + ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", data_size = ", data_size); + + // Encode all values into `buffer`. + const auto buffer = EncodeTestArray(data, bit_width); + + // Interleave Advance and GetBatch, verifying decoded values against the original. + RleBitPackedDecoder decoder(buffer.data(), static_cast(buffer.size()), + bit_width); + rle_size_t pos = 0; + auto advance = [&](rle_size_t* pos, rle_size_t n) { + n = std::min(n, data_size - *pos); + EXPECT_EQ(decoder.Advance(n), n); + *pos += n; + }; + auto read = [&](rle_size_t* pos, rle_size_t n) { + n = std::min(n, data_size - *pos); + std::vector got(n); + EXPECT_EQ(decoder.GetBatch(got.data(), n), n); + for (rle_size_t i = 0; i < n; ++i) { + EXPECT_EQ(got[i], data_values[*pos + i]) << "at position " << (*pos + i); + } + *pos += n; + }; + + // Advance in a `step` small relative to the data size, so we span many + // iterations and repeatedly cross run boundaries. + const rle_size_t step = std::max(data_size / 16, 1); + int iter = 0; + while (pos < data_size) { + // Some way to make various successions of `read` of `advance` + if (bit_width % 2 == iter % 3) { + read(&pos, step); + } + advance(&pos, step); + ++iter; + } + // Note: we do not assert exhaustion here because the encoder pads the last + // literal group to a multiple of 8 with zeros, leaving up to 7 extra values. +} + template struct DataTestRleBitPackedRandomPart { using value_type = T; @@ -1257,6 +1326,10 @@ void DoTestGetBatchSpacedRoundtrip() { CheckRoundTrip(*array, case_.bit_width, /* spaced= */ false, /* parts= */ 3); + // Tests for Advance + CheckAdvance(*array, case_.bit_width); + CheckAdvance(*array->Slice(1), case_.bit_width); + // Tests for GetBatchSpaced CheckRoundTrip(*array, case_.bit_width, /* spaced= */ true, /* parts= */ 1); From 2f060ad5523393c740c33db9e26d95423575b67f Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 09:18:20 +0200 Subject: [PATCH 06/13] Add RleBitPackedToBitmapDecoder::Advance --- cpp/src/arrow/util/rle_bitmap_internal.h | 76 ++++++++++----- cpp/src/arrow/util/rle_bitmap_test.cc | 118 +++++++++++++++++++++++ 2 files changed, 170 insertions(+), 24 deletions(-) diff --git a/cpp/src/arrow/util/rle_bitmap_internal.h b/cpp/src/arrow/util/rle_bitmap_internal.h index 0667a9f0a8ec..97d82804bd1c 100644 --- a/cpp/src/arrow/util/rle_bitmap_internal.h +++ b/cpp/src/arrow/util/rle_bitmap_internal.h @@ -287,6 +287,10 @@ class RleBitPackedToBitmapDecoder { /// values left or if an error occurred. [[nodiscard]] rle_size_t GetBatch(BitmapSpanMut out, rle_size_t batch_size); + /// Advance by as many values as provided or until exhaustion of the decoder. + /// Return the number of values skipped. + [[nodiscard]] rle_size_t Advance(rle_size_t batch_size); + private: RleBitPackedParser parser_ = {}; std::variant decoder_ = {}; @@ -296,10 +300,14 @@ class RleBitPackedToBitmapDecoder { return std::visit([](const auto& dec) { return dec.remaining(); }, decoder_); } - /// Get a batch of values from the current run and return the number elements read. - [[nodiscard]] rle_size_t RunGetBatch(BitmapSpanMut out, rle_size_t batch_size) { - return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder_); - } + /// Process data in the current run and subsequent ones. + /// + /// ``func`` is called as ``func(decoder, run_batch_size)`` where ``decoder`` is + /// the statically-typed run decoder and must return the number of values processed. + /// + /// Return the number of values processed. + template + rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size); }; /************************************************ @@ -320,50 +328,70 @@ struct RleBitPackedToBitmapDecoderGetDecoder { using type = BitPackedRunToBitmapDecoder; }; -inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, - rle_size_t batch_size) -> rle_size_t { +template +auto RleBitPackedToBitmapDecoder::ProcessValues(Callable&& func, + rle_size_t batch_size) -> rle_size_t { using ControlFlow = RleBitPackedParser::ControlFlow; if (ARROW_PREDICT_FALSE(batch_size == 0 || exhausted())) { return 0; } - rle_size_t values_read = 0; + rle_size_t values_processed = 0; // Remaining from a previous call that would have left some unread data from a run. if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { - const auto read = RunGetBatch(out, batch_size); - values_read += read; + const auto processed = + std::visit([&](auto& decoder) { return func(decoder, batch_size); }, decoder_); + values_processed += processed; // Either we fulfilled all the batch to be read or we finished remaining run. - if (ARROW_PREDICT_FALSE(values_read == batch_size)) { - return values_read; + if (ARROW_PREDICT_FALSE(values_processed == batch_size)) { + return values_processed; } ARROW_DCHECK(run_remaining() == 0); } parser_.ParseWithCallable([&](auto run) { - using RunDecoder = RleBitPackedToBitmapDecoderGetDecoder::type; + using RunDecoder = + typename RleBitPackedToBitmapDecoderGetDecoder::type; - ARROW_DCHECK_LT(values_read, batch_size); + ARROW_DCHECK_LT(values_processed, batch_size); + // Decode from a local decoder and only store it back into `decoder_` when we stop in + // the middle of a run. This keeps the run type statically known inside `func` and + // avoids writing the variant for every fully-consumed run. RunDecoder decoder(run); - // The output span carries its own bit offset, so advancing it past the values - // already written keeps successive runs correctly aligned in the bitmap. - const auto read = - decoder.GetBatch(out.NewStartingAt(values_read), batch_size - values_read); - ARROW_DCHECK_LE(read, batch_size - values_read); - values_read += read; - - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(values_read == batch_size || read == 0)) { + const auto read = func(decoder, batch_size - values_processed); + ARROW_DCHECK_LE(read, batch_size - values_processed); + values_processed += read; + + if (ARROW_PREDICT_FALSE(values_processed == batch_size || read == 0)) { decoder_ = std::move(decoder); return ControlFlow::Break; } - return ControlFlow::Continue; }); - return values_read; + return values_processed; +} + +inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, + rle_size_t batch_size) -> rle_size_t { + return ProcessValues( + [out](auto& decoder, rle_size_t run_batch_size) mutable { + const auto read = decoder.GetBatch(out, run_batch_size); + out = out.NewStartingAt(read); + return read; + }, + batch_size); +} + +inline auto RleBitPackedToBitmapDecoder::Advance(rle_size_t batch_size) -> rle_size_t { + return ProcessValues( + [](auto& decoder, rle_size_t run_batch_size) { + return decoder.Advance(run_batch_size); + }, + batch_size); } } // namespace arrow::util diff --git a/cpp/src/arrow/util/rle_bitmap_test.cc b/cpp/src/arrow/util/rle_bitmap_test.cc index 7dc572eb5ad7..79762d947c2b 100644 --- a/cpp/src/arrow/util/rle_bitmap_test.cc +++ b/cpp/src/arrow/util/rle_bitmap_test.cc @@ -468,6 +468,83 @@ void CheckRleBitPackedToBitmap(const std::vector& bytes, } } +/// Interleave Advance and GetBatch across the whole stream with a fixed `step`, +/// verifying the values that are read. +/// +/// Advance and GetBatch both go through the parser, so `step` values that do not +/// align with run boundaries leave the decoder mid-run between calls, exercising +/// the resume path they share. +void CheckRleBitPackedAdvance(const std::vector& bytes, + const std::vector& expected, rle_size_t step) { + ARROW_SCOPED_TRACE("step = ", step); + const auto n_vals = static_cast(expected.size()); + + // Advancing over the whole stream consumes every value then reports exhaustion. + { + RleBitPackedToBitmapDecoder decoder(bytes.data(), + static_cast(bytes.size())); + rle_size_t pos = 0; + while (pos < n_vals) { + const auto n = std::min(step, n_vals - pos); + EXPECT_EQ(decoder.Advance(n), n) << "at pos " << pos; + pos += n; + } + EXPECT_EQ(pos, n_vals); + EXPECT_TRUE(decoder.exhausted()); + EXPECT_EQ(decoder.Advance(1), 0); + } + + // Interleave Advance and GetBatch, checking the values we do read. + { + RleBitPackedToBitmapDecoder decoder(bytes.data(), + static_cast(bytes.size())); + // Output buffer with one guard byte to catch out-of-bounds writes. + std::vector out(static_cast(bit_util::BytesForBits(n_vals)) + 1, 0); + + rle_size_t pos = 0; + auto advance = [&](rle_size_t n) { + n = std::min(n, n_vals - pos); + EXPECT_EQ(decoder.Advance(n), n) << "at pos " << pos; + pos += n; + }; + auto read = [&](rle_size_t n) { + n = std::min(n, n_vals - pos); + const auto got = decoder.GetBatch(BitmapSpanMut(out.data(), /*bit_start=*/pos), n); + EXPECT_EQ(got, n) << "at pos " << pos; + CheckDecodedBits({ + .actual = out, + .expected = expected, + .count = n, + .actual_start_bit = pos, + .expected_start_idx = pos, + }); + pos += n; + }; + + int iter = 0; + while (pos < n_vals) { + // Vary how reads and advances alternate so runs are consumed many ways. + if (iter % 3 != 0) { + read(step); + } + advance(step); + ++iter; + } + EXPECT_TRUE(decoder.exhausted()); + } +} + +/// Run the Advance check over a battery of step sizes. +void CheckRleBitPackedToBitmapAdvance(const std::vector& bytes, + const std::vector& expected) { + const auto n_vals = static_cast(expected.size()); + ASSERT_GT(n_vals, 0); + for (const rle_size_t step : {rle_size_t{1}, rle_size_t{3}, rle_size_t{7}, + rle_size_t{8}, rle_size_t{9}, rle_size_t{33}}) { + CheckRleBitPackedAdvance(bytes, expected, step); + } +} + } // namespace TEST(RleBitPackedToBitmapDecoder, Empty) { @@ -599,4 +676,45 @@ TEST(RleBitPackedToBitmapDecoder, Truncated) { CheckDecodedBits({.actual = out, .expected = expected, .count = 10}); } +TEST(RleBitPackedToBitmapDecoder, AdvanceSingleRun) { + { + std::vector bytes; + std::vector expected; + AppendRleRun(bytes, expected, /*value=*/true, /*count=*/100); + CheckRleBitPackedToBitmapAdvance(bytes, expected); + } + { + std::vector bytes; + std::vector expected; + AppendBitPackedRun(bytes, expected, {0b10101010, 0b11001100, 0b11110000}); + CheckRleBitPackedToBitmapAdvance(bytes, expected); + } +} + +TEST(RleBitPackedToBitmapDecoder, AdvanceMixedRunsAligned) { + // All runs end on a byte boundary. + std::vector bytes; + std::vector expected; + AppendRleRun(bytes, expected, /*value=*/false, /*count=*/16); + AppendBitPackedRun(bytes, expected, {0b10101010, 0b01010101}); + AppendRleRun(bytes, expected, /*value=*/true, /*count=*/64); + AppendBitPackedRun(bytes, expected, {0b00001111}); + CheckRleBitPackedToBitmapAdvance(bytes, expected); +} + +TEST(RleBitPackedToBitmapDecoder, AdvanceMixedRunsUnaligned) { + // RLE runs with counts that are not multiples of 8 make each following run + // start at a non-byte-aligned output position. + std::vector bytes; + std::vector expected; + AppendRleRun(bytes, expected, /*value=*/true, /*count=*/13); + AppendBitPackedRun(bytes, expected, {0b01101001, 0b10010110}); + AppendRleRun(bytes, expected, /*value=*/false, /*count=*/5); + AppendRleRun(bytes, expected, /*value=*/true, /*count=*/200); + AppendBitPackedRun(bytes, expected, {0b11110000}); + AppendRleRun(bytes, expected, /*value=*/false, /*count=*/3); + AppendBitPackedRun(bytes, expected, {0b10110001, 0b00011101}); + CheckRleBitPackedToBitmapAdvance(bytes, expected); +} + } // namespace arrow::util From e31ca0dabe8584515f3bdcea6495b49853897ab7 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 09:37:23 +0200 Subject: [PATCH 07/13] Add legacy BitPackedDecoder --- cpp/src/arrow/util/rle_encoding_internal.h | 56 +++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 11f36c498e23..db0564b76630 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -157,7 +157,8 @@ class BitPackedRun { constexpr BitPackedRun() noexcept = default; constexpr BitPackedRun(const uint8_t* data, rle_size_t values_count, - rle_size_t value_bit_width, rle_size_t max_read_bytes) noexcept + rle_size_t value_bit_width, + rle_size_t max_read_bytes = -1) noexcept : data_(data), values_count_(values_count), max_read_bytes_(max_read_bytes) { ARROW_DCHECK_GE(value_bit_width, 0); ARROW_DCHECK_GE(values_count_, 0); @@ -522,6 +523,59 @@ class RleBitPackedDecoder { int64_t valid_bits_offset, rle_size_t null_count); }; +/// Minimal decoder for legacy bit packed encoding (BIT_PACKED = 4). +/// +/// The number of values that the decoder can represent is up to 2^31 - 1. +template +class BitPackedDecoder : private BitPackedRunDecoder { + private: + using Base = BitPackedRunDecoder; + + public: + /// The type in which the data should be decoded. + using value_type = T; + + using Base::Advance; + + BitPackedDecoder() noexcept = default; + + /// Create a decoder object. + /// + /// data and data_size are the raw bytes to decode. + /// value_bit_width is the size in bits of each encoded value. + BitPackedDecoder(const uint8_t* data, rle_size_t data_size, + rle_size_t value_bit_width) noexcept { + Reset(data, data_size, value_bit_width); + } + + void Reset(const uint8_t* data, rle_size_t data_size, + rle_size_t value_bit_width) noexcept { + value_bit_width_ = value_bit_width; + const auto value_count = (static_cast(data_size) * 8) / value_bit_width; + ARROW_DCHECK_LE(value_count, std::numeric_limits::max()); + const auto run = BitPackedRun{ + /* data= */ data, + /* value_count= */ static_cast(value_count), + /* value_bit_width= */ value_bit_width, + }; + return Base::Reset(run, value_bit_width); + } + + /// Whether there is still values to iterate over. + bool exhausted() const { return Base::remaining() == 0; } + + /// 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_); } + + /// 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_); + } + + private: + rle_size_t value_bit_width_ = {}; +}; + /// Class to incrementally build the rle data. This class does not allocate any memory. /// The encoding has two modes: encoding repeated runs and literal runs. /// If the run is sufficiently short, it is more efficient to encode as a literal run. From d6f43d48c2c7339ebf328c604a4fd853b69c6edb Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 25 Jun 2026 11:38:51 +0200 Subject: [PATCH 08/13] Simplify LevelDecoder::Impl --- cpp/src/parquet/column_reader.cc | 57 +------------------------------- 1 file changed, 1 insertion(+), 56 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index ad5a9ee74968..6e565283288c 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -100,67 +100,12 @@ constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = * LevelDecoder * ******************/ -namespace { - -/// Adapter around ``BitReader`` that mimics the API of ``RleBitPackedDecoder``. -/// -/// Best would be to make this a first class citizen, possibly reusing -/// ``BitPackedRunDecoder`` but the preconditions on the possible number of values -/// that the run can represent differ. -template -class BitPackedDecoderWrapper : private ::arrow::bit_util::BitReader { - public: - /// The type in which the data should be decoded. - using value_type = T; - using rle_size_t = ::arrow::util::rle_size_t; - - BitPackedDecoderWrapper() noexcept = default; - - BitPackedDecoderWrapper(const uint8_t* data, rle_size_t data_size, - rle_size_t value_bit_width) noexcept { - Reset(data, data_size, value_bit_width); - } - - void Reset(const uint8_t* data, rle_size_t data_size, - rle_size_t value_bit_width) noexcept { - value_bit_width_ = value_bit_width; - return Base::Reset(data, data_size); - } - - /// Whether there is still values to iterate over. - bool exhausted() const { return Base::bytes_left() >= value_bit_width_; } - - /// Gets the next value or returns false if there are no more or an error occurred. - [[nodiscard]] bool Get(value_type* val) { - return Base::GetValue(value_bit_width_, val); - } - - /// 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(value_bit_width_, out, batch_size); - } - - private: - using Base = ::arrow::bit_util::BitReader; - - rle_size_t value_bit_width_ = {}; -}; -} // namespace - struct LevelDecoder::Impl { using RleBitPackedDecoder = ::arrow::util::RleBitPackedDecoder; - using BitPackedDecoder = BitPackedDecoderWrapper; + using BitPackedDecoder = ::arrow::util::BitPackedDecoder; std::variant decoder = {}; - constexpr bool is_rle_bit_packed() const { - return std::holds_alternative(decoder); - } - - constexpr bool is_bit_packed() const { - return std::holds_alternative(decoder); - } - [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); } From 06464db83c6e302f6cb6b8ad36dfd64ee7db9e39 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 09:48:01 +0200 Subject: [PATCH 09/13] Add LevelDecoder::Skip --- cpp/src/parquet/column_reader.cc | 12 ++++++++++++ cpp/src/parquet/column_reader.h | 7 ++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 6e565283288c..ef732a91abe2 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -109,6 +109,10 @@ struct LevelDecoder::Impl { [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); } + + [[nodiscard]] int Advance(int batch_size) { + return std::visit([&](auto& dec) { return dec.Advance(batch_size); }, decoder); + } }; LevelDecoder::LevelDecoder(int16_t max_level) @@ -192,6 +196,14 @@ int LevelDecoder::Decode(int batch_size, int16_t* levels) { return num_decoded; } +int LevelDecoder::Skip(int batch_size) { + const int num_values = std::min(num_values_remaining_, batch_size); + const int num_advanced = decoder_->Advance(num_values); + ARROW_DCHECK_EQ(num_values, num_advanced); + num_values_remaining_ -= num_advanced; + return num_advanced; +} + ReaderProperties default_reader_properties() { static ReaderProperties default_reader_properties; return default_reader_properties; diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index 802e0062cc2f..fe5efa779d8d 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -81,9 +81,12 @@ class PARQUET_EXPORT LevelDecoder { void SetDataV2(int32_t num_bytes, int16_t max_level, int num_buffered_values, const uint8_t* data); - /// Decode a batch of levels into an array and returns the number of levels decoded + /// Decode a batch of levels into an array and returns the number of levels decoded. int Decode(int batch_size, int16_t* levels); + /// Advance the decoder and throw away decoder levels. + int Skip(int batch_size); + /// Return the max level used in this decoder. int max_level() const { return max_level_; } @@ -91,6 +94,8 @@ class PARQUET_EXPORT LevelDecoder { struct Impl; std::unique_ptr decoder_; + /// Number of value remaining. The underlying decoder zero pads bit packed values + /// up to a multiple of 8 so it cannot know the exact number of remaining values. int num_values_remaining_ = 0; int16_t max_level_; }; From 481cf9ab7db9f9e9aa32439e30fe982473a3da77 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 14:48:19 +0200 Subject: [PATCH 10/13] Add RleBitPackedDecoder::CountUpTo --- cpp/src/arrow/util/rle_encoding_internal.h | 80 ++++++++++++++++++++++ cpp/src/arrow/util/rle_encoding_test.cc | 25 +++++++ 2 files changed, 105 insertions(+) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index db0564b76630..efedffaf61e9 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -259,6 +259,18 @@ class RleBitPackedParser { std::pair PeekImpl(Handler&&) const; }; +template +struct RleCountUpToParams { + T value; + rle_size_t batch_size; + rle_size_t value_bit_width; +}; + +struct RleCountUpToResult { + rle_size_t count; + rle_size_t advanced_count; +}; + /// Decoder class for a single run of RLE encoded data. template class RleRunDecoder { @@ -301,6 +313,15 @@ class RleRunDecoder { return steps; } + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(const RleCountUpToParams& p) { + const auto steps = Advance(p.batch_size); + return {.count = steps * (p.value == value_), .advanced_count = steps}; + } + /// 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; @@ -364,6 +385,29 @@ class BitPackedRunDecoder { return steps; } + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(const RleCountUpToParams& p) { + // Decoding in a stack buffer of 512 values: 1KB for int16_t used in levels + // and up to 4KB for uint64_t. + constexpr rle_size_t kBufferValueCount = 512; + alignas(16) value_type buffer[kBufferValueCount]; // uninitialized + + rle_size_t remaining = p.batch_size; + rle_size_t count = 0; + rle_size_t read_iter = 0; + do { + const auto batch_iter = std::min(remaining, kBufferValueCount); + read_iter = GetBatch(buffer, batch_iter, p.value_bit_width); + count += static_cast(std::count(buffer, buffer + read_iter, p.value)); + remaining -= read_iter; + } while (remaining > 0 && read_iter > 0); + + return {.count = count, .advanced_count = p.batch_size - remaining}; + } + /// 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; @@ -453,6 +497,12 @@ class RleBitPackedDecoder { /// Return the number of values skipped. [[nodiscard]] rle_size_t Advance(rle_size_t batch_size); + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(value_type value, rle_size_t batch_size); + /// Gets the next value or returns false if there are no more or an error occurred. /// /// NB: Because the encoding only supports literal runs with lengths @@ -564,6 +614,18 @@ class BitPackedDecoder : private BitPackedRunDecoder { /// Whether there is still values to iterate over. bool exhausted() const { return Base::remaining() == 0; } + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(value_type value, rle_size_t batch_size) { + return Base::CountUpTo({ + .value = value, + .batch_size = batch_size, + .value_bit_width = value_bit_width_, + }); + } + /// 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_); } @@ -897,6 +959,24 @@ auto RleBitPackedDecoder::Advance(rle_size_t batch_size) -> rle_size_t { batch_size); } +template +RleCountUpToResult RleBitPackedDecoder::CountUpTo(value_type value, + rle_size_t batch_size) { + rle_size_t count = 0; + const rle_size_t advanced_count = ProcessValues( + [value, this, &count](auto& decoder, rle_size_t run_batch_size) { + const auto result = decoder.CountUpTo({ + .value = value, + .batch_size = run_batch_size, + .value_bit_width = value_bit_width_, + }); + count += result.count; + return result.advanced_count; + }, + batch_size); + return {.count = count, .advanced_count = advanced_count}; +} + template bool RleBitPackedDecoder::Get(value_type* val) { return GetBatch(val, 1) == 1; diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index a371f7514441..3150946797f9 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -335,6 +335,31 @@ TEST(Rle, RleDecoder) { /* expected_value= */ 16777749); } +TEST(Rle, RleDecoderCountUpTo) { + // A run of value 21 repeated 23 times. + const std::array bytes = {21, 0, 0}; + const auto run = RleRun(bytes.data(), /* value_count= */ 23, /* bit_width= */ 5); + auto decoder = RleRunDecoder(run, /* value_bit_width= */ 5); + + // Counting the repeated value counts every advanced element. + auto res = decoder.CountUpTo({.value = 21, .batch_size = 10, .value_bit_width = 5}); + EXPECT_EQ(res.advanced_count, 10); + EXPECT_EQ(res.count, 10); + EXPECT_EQ(decoder.remaining(), 23 - 10); + + // Counting another value matches nothing but still advances. + res = decoder.CountUpTo({.value = 99, .batch_size = 5, .value_bit_width = 5}); + EXPECT_EQ(res.advanced_count, 5); + EXPECT_EQ(res.count, 0); + EXPECT_EQ(decoder.remaining(), 23 - 15); + + // Requesting more than remaining is capped to the remaining count. + res = decoder.CountUpTo({.value = 21, .batch_size = 100, .value_bit_width = 5}); + EXPECT_EQ(res.advanced_count, 8); + EXPECT_EQ(res.count, 8); + EXPECT_EQ(decoder.remaining(), 0); +} + template void TestBitPackedDecoder(std::vector bytes, rle_size_t value_count, rle_size_t bit_width, std::vector expected) { From bd244ad58cd8d905be450ef123885567a2b2d7d9 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 15:04:02 +0200 Subject: [PATCH 11/13] Add LevelDecoder::CountUpTo --- cpp/src/parquet/column_reader.cc | 13 +++++++++++++ cpp/src/parquet/column_reader.h | 6 ++++++ 2 files changed, 19 insertions(+) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index ef732a91abe2..3c66cf82558c 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -113,6 +113,11 @@ struct LevelDecoder::Impl { [[nodiscard]] int Advance(int batch_size) { return std::visit([&](auto& dec) { return dec.Advance(batch_size); }, decoder); } + + auto CountUpTo(int16_t value, int batch_size) { + return std::visit([&](auto& dec) { return dec.CountUpTo(value, batch_size); }, + decoder); + } }; LevelDecoder::LevelDecoder(int16_t max_level) @@ -204,6 +209,14 @@ int LevelDecoder::Skip(int batch_size) { return num_advanced; } +std::pair LevelDecoder::CountUpTo(int16_t value, int batch_size) { + const int num_values = std::min(num_values_remaining_, batch_size); + const auto result = decoder_->CountUpTo(value, num_values); + ARROW_DCHECK_EQ(num_values, result.advanced_count); + num_values_remaining_ -= result.advanced_count; + return {result.count, result.advanced_count}; +} + ReaderProperties default_reader_properties() { static ReaderProperties default_reader_properties; return default_reader_properties; diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index fe5efa779d8d..33d6d37b49b6 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -87,6 +87,12 @@ class PARQUET_EXPORT LevelDecoder { /// Advance the decoder and throw away decoder levels. int Skip(int batch_size); + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of elements that were processed. + std::pair CountUpTo(int16_t value, int batch_size); + /// Return the max level used in this decoder. int max_level() const { return max_level_; } From cbbe6d6e2390e328a9e4d6138a42daf76ef2b0d2 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 16:05:00 +0200 Subject: [PATCH 12/13] Remove scratch_for_skip_ --- cpp/src/parquet/column_reader.cc | 98 ++++++++++++++------------------ 1 file changed, 42 insertions(+), 56 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 3c66cf82558c..596fc9b44bc8 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -681,6 +681,12 @@ class SkippableTypedDecoder { Decoder* get() { return decoder_; } + const Decoder* operator->() const { return decoder_; } + + Decoder* operator->() { return decoder_; } + + explicit operator bool() const { return decoder_ != nullptr; } + int64_t Skip(int64_t num_values, ::arrow::MemoryPool* pool = nullptr) { EnsureScratch(pool); @@ -846,7 +852,7 @@ class ColumnReaderImplBase { } new_dictionary_ = true; - current_decoder_ = decoders_[encoding].get(); + current_decoder_.SetDecoder(decoders_[encoding].get()); ARROW_DCHECK(current_decoder_); } @@ -949,7 +955,7 @@ class ColumnReaderImplBase { auto it = decoders_.find(static_cast(encoding)); if (it != decoders_.end()) { ARROW_DCHECK(it->second.get() != nullptr); - current_decoder_ = it->second.get(); + current_decoder_.SetDecoder(it->second.get()); } else { switch (encoding) { case Encoding::PLAIN: @@ -959,7 +965,7 @@ class ColumnReaderImplBase { case Encoding::DELTA_BYTE_ARRAY: case Encoding::DELTA_LENGTH_BYTE_ARRAY: { auto decoder = MakeTypedDecoder(encoding, descr_, pool_); - current_decoder_ = decoder.get(); + current_decoder_.SetDecoder(decoder.get()); decoders_[static_cast(encoding)] = std::move(decoder); break; } @@ -1018,7 +1024,7 @@ class ColumnReaderImplBase { ::arrow::MemoryPool* pool_; using DecoderType = TypedDecoder; - DecoderType* current_decoder_ = nullptr; + SkippableTypedDecoder current_decoder_; Encoding::type current_encoding_ = Encoding::UNKNOWN; /// Flag to signal when a new dictionary has been set, for the benefit of @@ -1074,19 +1080,11 @@ class TypedColumnReaderImpl : public TypedColumnReader, this->exposed_encoding_ = encoding; } - // Allocate enough scratch space to accommodate skipping 16-bit levels or any - // value type. - void InitScratchForSkip(); - - // Scratch space for reading and throwing away rep/def levels and values when - // skipping. - std::shared_ptr scratch_for_skip_; - private: // Read dictionary indices. Similar to ReadValues but decode data to dictionary indices. // This function is called only by ReadBatchWithDictionary(). int64_t ReadDictionaryIndices(int64_t indices_to_read, int32_t* indices) { - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); return decoder->DecodeIndices(static_cast(indices_to_read), indices); } @@ -1094,7 +1092,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, // owned by the internal decoder and is destroyed when the reader is destroyed. This // function is called only by ReadBatchWithDictionary() after dictionary is configured. void GetDictionary(const T** dictionary, int32_t* dictionary_length) { - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); decoder->GetDictionary(dictionary, dictionary_length); } @@ -1225,40 +1223,40 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size, int16_t* def return total_values; } -template -void TypedColumnReaderImpl::InitScratchForSkip() { - if (this->scratch_for_skip_ == nullptr) { - int value_size = type_traits::value_byte_size; - this->scratch_for_skip_ = AllocateBuffer( - this->pool_, kSkipScratchBatchSize * std::max(sizeof(int16_t), value_size)); - } -} - template int64_t TypedColumnReaderImpl::Skip(int64_t num_values_to_skip) { int64_t values_to_skip = num_values_to_skip; // Optimization: Do not call HasNext() when values_to_skip == 0. while (values_to_skip > 0 && HasNext()) { // If the number of values to skip is more than the number of undecoded values, skip - // the Page. + // the whole Page without decoding levels or values. const int64_t available_values = this->available_values_current_page(); if (values_to_skip >= available_values) { values_to_skip -= available_values; this->ConsumeBufferedValues(available_values); } else { - // We need to read this Page - // Jump to the right offset in the Page - int64_t values_read = 0; - InitScratchForSkip(); - ARROW_DCHECK_NE(this->scratch_for_skip_, nullptr); - do { - int64_t batch_size = std::min(kSkipScratchBatchSize, values_to_skip); - values_read = ReadBatch(static_cast(batch_size), - scratch_for_skip_->mutable_data_as(), - scratch_for_skip_->mutable_data_as(), - scratch_for_skip_->mutable_data_as(), &values_read); - values_to_skip -= values_read; - } while (values_read > 0 && values_to_skip > 0); + // Skip within the current Page. Since `values_to_skip < available_values`, the + // whole batch fits inside this Page and no page boundary is crossed. + const int batch_size = static_cast(values_to_skip); + + // Advance the definition levels, counting how many correspond to present + // (non-null) values that must be skipped in the data decoder. + int64_t non_null_values_to_skip = batch_size; + if (this->max_def_level() > 0) { + const auto [count, advanced] = + this->definition_level_decoder_.CountUpTo(this->max_def_level(), batch_size); + non_null_values_to_skip = count; + ARROW_DCHECK_EQ(advanced, batch_size); + } + // Advance the repetition levels; their values are not needed. + if (this->max_rep_level() > 0) { + this->repetition_level_decoder_.Skip(batch_size); + } + // Skip the corresponding data values. + this->current_decoder_.Skip(non_null_values_to_skip, this->pool_); + + this->ConsumeBufferedValues(batch_size); + values_to_skip -= batch_size; } } return num_values_to_skip - values_to_skip; @@ -1355,7 +1353,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } const void* ReadDictionary(int32_t* dictionary_length) override { - if (this->current_decoder_ == nullptr && !this->HasNextInternal()) { + if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; return nullptr; } @@ -1368,7 +1366,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, << EncodingToString(this->current_encoding_); throw ParquetException(ss.str()); } - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); return reinterpret_cast(dictionary); @@ -1604,20 +1602,8 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Read 'num_values' values and throw them away. // Throws an error if it could not read 'num_values'. void ReadAndThrowAwayValues(int64_t num_values) { - int64_t values_left = num_values; - int64_t values_read = 0; - - // Allocate enough scratch space to accommodate 16-bit levels or any - // value type - this->InitScratchForSkip(); - ARROW_DCHECK_NE(this->scratch_for_skip_, nullptr); - do { - int64_t batch_size = std::min(kSkipScratchBatchSize, values_left); - values_read = this->ReadValues( - batch_size, this->scratch_for_skip_->template mutable_data_as()); - values_left -= values_read; - } while (values_read > 0 && values_left > 0); - if (values_left > 0) { + const int64_t values_read = this->current_decoder_.Skip(num_values, this->pool_); + if (values_read < num_values) { std::stringstream ss; ss << "Could not read and throw away " << num_values << " values"; throw ParquetException(ss.str()); @@ -2216,7 +2202,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); decoder->InsertDictionary(&builder_); this->new_dictionary_ = false; } @@ -2226,7 +2212,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); num_decoded = decoder->DecodeIndices(static_cast(values_to_read), &builder_); } else { num_decoded = this->current_decoder_->DecodeArrowNonNull( @@ -2241,7 +2227,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); num_decoded = decoder->DecodeIndicesSpaced( static_cast(values_to_read), static_cast(null_count), valid_bits_->mutable_data(), values_written_, &builder_); From 34c1b153c8cae6b18fef9faad83db3d3eec5fd10 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 8 Jul 2026 16:58:15 +0200 Subject: [PATCH 13/13] Remove RleBiPackedToBitmap changes --- cpp/src/arrow/util/rle_bitmap_internal.h | 76 +++++---------- cpp/src/arrow/util/rle_bitmap_test.cc | 118 ----------------------- 2 files changed, 24 insertions(+), 170 deletions(-) diff --git a/cpp/src/arrow/util/rle_bitmap_internal.h b/cpp/src/arrow/util/rle_bitmap_internal.h index 97d82804bd1c..0667a9f0a8ec 100644 --- a/cpp/src/arrow/util/rle_bitmap_internal.h +++ b/cpp/src/arrow/util/rle_bitmap_internal.h @@ -287,10 +287,6 @@ class RleBitPackedToBitmapDecoder { /// values left or if an error occurred. [[nodiscard]] rle_size_t GetBatch(BitmapSpanMut out, rle_size_t batch_size); - /// Advance by as many values as provided or until exhaustion of the decoder. - /// Return the number of values skipped. - [[nodiscard]] rle_size_t Advance(rle_size_t batch_size); - private: RleBitPackedParser parser_ = {}; std::variant decoder_ = {}; @@ -300,14 +296,10 @@ class RleBitPackedToBitmapDecoder { return std::visit([](const auto& dec) { return dec.remaining(); }, decoder_); } - /// Process data in the current run and subsequent ones. - /// - /// ``func`` is called as ``func(decoder, run_batch_size)`` where ``decoder`` is - /// the statically-typed run decoder and must return the number of values processed. - /// - /// Return the number of values processed. - template - rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size); + /// Get a batch of values from the current run and return the number elements read. + [[nodiscard]] rle_size_t RunGetBatch(BitmapSpanMut out, rle_size_t batch_size) { + return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder_); + } }; /************************************************ @@ -328,70 +320,50 @@ struct RleBitPackedToBitmapDecoderGetDecoder { using type = BitPackedRunToBitmapDecoder; }; -template -auto RleBitPackedToBitmapDecoder::ProcessValues(Callable&& func, - rle_size_t batch_size) -> rle_size_t { +inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, + rle_size_t batch_size) -> rle_size_t { using ControlFlow = RleBitPackedParser::ControlFlow; if (ARROW_PREDICT_FALSE(batch_size == 0 || exhausted())) { return 0; } - rle_size_t values_processed = 0; + rle_size_t values_read = 0; // Remaining from a previous call that would have left some unread data from a run. if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { - const auto processed = - std::visit([&](auto& decoder) { return func(decoder, batch_size); }, decoder_); - values_processed += processed; + const auto read = RunGetBatch(out, batch_size); + values_read += read; // Either we fulfilled all the batch to be read or we finished remaining run. - if (ARROW_PREDICT_FALSE(values_processed == batch_size)) { - return values_processed; + if (ARROW_PREDICT_FALSE(values_read == batch_size)) { + return values_read; } ARROW_DCHECK(run_remaining() == 0); } parser_.ParseWithCallable([&](auto run) { - using RunDecoder = - typename RleBitPackedToBitmapDecoderGetDecoder::type; + using RunDecoder = RleBitPackedToBitmapDecoderGetDecoder::type; - ARROW_DCHECK_LT(values_processed, batch_size); - // Decode from a local decoder and only store it back into `decoder_` when we stop in - // the middle of a run. This keeps the run type statically known inside `func` and - // avoids writing the variant for every fully-consumed run. + ARROW_DCHECK_LT(values_read, batch_size); RunDecoder decoder(run); - const auto read = func(decoder, batch_size - values_processed); - ARROW_DCHECK_LE(read, batch_size - values_processed); - values_processed += read; - - if (ARROW_PREDICT_FALSE(values_processed == batch_size || read == 0)) { + // The output span carries its own bit offset, so advancing it past the values + // already written keeps successive runs correctly aligned in the bitmap. + const auto read = + decoder.GetBatch(out.NewStartingAt(values_read), batch_size - values_read); + ARROW_DCHECK_LE(read, batch_size - values_read); + values_read += read; + + // Stop reading and store remaining decoder + if (ARROW_PREDICT_FALSE(values_read == batch_size || read == 0)) { decoder_ = std::move(decoder); return ControlFlow::Break; } + return ControlFlow::Continue; }); - return values_processed; -} - -inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, - rle_size_t batch_size) -> rle_size_t { - return ProcessValues( - [out](auto& decoder, rle_size_t run_batch_size) mutable { - const auto read = decoder.GetBatch(out, run_batch_size); - out = out.NewStartingAt(read); - return read; - }, - batch_size); -} - -inline auto RleBitPackedToBitmapDecoder::Advance(rle_size_t batch_size) -> rle_size_t { - return ProcessValues( - [](auto& decoder, rle_size_t run_batch_size) { - return decoder.Advance(run_batch_size); - }, - batch_size); + return values_read; } } // namespace arrow::util diff --git a/cpp/src/arrow/util/rle_bitmap_test.cc b/cpp/src/arrow/util/rle_bitmap_test.cc index 79762d947c2b..7dc572eb5ad7 100644 --- a/cpp/src/arrow/util/rle_bitmap_test.cc +++ b/cpp/src/arrow/util/rle_bitmap_test.cc @@ -468,83 +468,6 @@ void CheckRleBitPackedToBitmap(const std::vector& bytes, } } -/// Interleave Advance and GetBatch across the whole stream with a fixed `step`, -/// verifying the values that are read. -/// -/// Advance and GetBatch both go through the parser, so `step` values that do not -/// align with run boundaries leave the decoder mid-run between calls, exercising -/// the resume path they share. -void CheckRleBitPackedAdvance(const std::vector& bytes, - const std::vector& expected, rle_size_t step) { - ARROW_SCOPED_TRACE("step = ", step); - const auto n_vals = static_cast(expected.size()); - - // Advancing over the whole stream consumes every value then reports exhaustion. - { - RleBitPackedToBitmapDecoder decoder(bytes.data(), - static_cast(bytes.size())); - rle_size_t pos = 0; - while (pos < n_vals) { - const auto n = std::min(step, n_vals - pos); - EXPECT_EQ(decoder.Advance(n), n) << "at pos " << pos; - pos += n; - } - EXPECT_EQ(pos, n_vals); - EXPECT_TRUE(decoder.exhausted()); - EXPECT_EQ(decoder.Advance(1), 0); - } - - // Interleave Advance and GetBatch, checking the values we do read. - { - RleBitPackedToBitmapDecoder decoder(bytes.data(), - static_cast(bytes.size())); - // Output buffer with one guard byte to catch out-of-bounds writes. - std::vector out(static_cast(bit_util::BytesForBits(n_vals)) + 1, 0); - - rle_size_t pos = 0; - auto advance = [&](rle_size_t n) { - n = std::min(n, n_vals - pos); - EXPECT_EQ(decoder.Advance(n), n) << "at pos " << pos; - pos += n; - }; - auto read = [&](rle_size_t n) { - n = std::min(n, n_vals - pos); - const auto got = decoder.GetBatch(BitmapSpanMut(out.data(), /*bit_start=*/pos), n); - EXPECT_EQ(got, n) << "at pos " << pos; - CheckDecodedBits({ - .actual = out, - .expected = expected, - .count = n, - .actual_start_bit = pos, - .expected_start_idx = pos, - }); - pos += n; - }; - - int iter = 0; - while (pos < n_vals) { - // Vary how reads and advances alternate so runs are consumed many ways. - if (iter % 3 != 0) { - read(step); - } - advance(step); - ++iter; - } - EXPECT_TRUE(decoder.exhausted()); - } -} - -/// Run the Advance check over a battery of step sizes. -void CheckRleBitPackedToBitmapAdvance(const std::vector& bytes, - const std::vector& expected) { - const auto n_vals = static_cast(expected.size()); - ASSERT_GT(n_vals, 0); - for (const rle_size_t step : {rle_size_t{1}, rle_size_t{3}, rle_size_t{7}, - rle_size_t{8}, rle_size_t{9}, rle_size_t{33}}) { - CheckRleBitPackedAdvance(bytes, expected, step); - } -} - } // namespace TEST(RleBitPackedToBitmapDecoder, Empty) { @@ -676,45 +599,4 @@ TEST(RleBitPackedToBitmapDecoder, Truncated) { CheckDecodedBits({.actual = out, .expected = expected, .count = 10}); } -TEST(RleBitPackedToBitmapDecoder, AdvanceSingleRun) { - { - std::vector bytes; - std::vector expected; - AppendRleRun(bytes, expected, /*value=*/true, /*count=*/100); - CheckRleBitPackedToBitmapAdvance(bytes, expected); - } - { - std::vector bytes; - std::vector expected; - AppendBitPackedRun(bytes, expected, {0b10101010, 0b11001100, 0b11110000}); - CheckRleBitPackedToBitmapAdvance(bytes, expected); - } -} - -TEST(RleBitPackedToBitmapDecoder, AdvanceMixedRunsAligned) { - // All runs end on a byte boundary. - std::vector bytes; - std::vector expected; - AppendRleRun(bytes, expected, /*value=*/false, /*count=*/16); - AppendBitPackedRun(bytes, expected, {0b10101010, 0b01010101}); - AppendRleRun(bytes, expected, /*value=*/true, /*count=*/64); - AppendBitPackedRun(bytes, expected, {0b00001111}); - CheckRleBitPackedToBitmapAdvance(bytes, expected); -} - -TEST(RleBitPackedToBitmapDecoder, AdvanceMixedRunsUnaligned) { - // RLE runs with counts that are not multiples of 8 make each following run - // start at a non-byte-aligned output position. - std::vector bytes; - std::vector expected; - AppendRleRun(bytes, expected, /*value=*/true, /*count=*/13); - AppendBitPackedRun(bytes, expected, {0b01101001, 0b10010110}); - AppendRleRun(bytes, expected, /*value=*/false, /*count=*/5); - AppendRleRun(bytes, expected, /*value=*/true, /*count=*/200); - AppendBitPackedRun(bytes, expected, {0b11110000}); - AppendRleRun(bytes, expected, /*value=*/false, /*count=*/3); - AppendBitPackedRun(bytes, expected, {0b10110001, 0b00011101}); - CheckRleBitPackedToBitmapAdvance(bytes, expected); -} - } // namespace arrow::util