From 5d2b73ccbf16c6c38b3f469a8ebd9637646e5e24 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 9 Jul 2026 11:32:37 +0200 Subject: [PATCH 01/33] Move TypedColReaderImpl::Skip --- cpp/src/parquet/column_reader.cc | 81 ++++++++++++++++---------------- 1 file changed, 41 insertions(+), 40 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 162a72bd157..84d95f05d0e 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1048,6 +1048,44 @@ class ColumnReaderImplBase { std::unordered_map> decoders_; void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } + + int64_t 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 && HasNextInternal()) { + // If the number of values to skip is more than the number of undecoded values, skip + // 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 { + // 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 = this->definition_level_decoder_.CountUpTo( + this->max_def_level(), batch_size); + non_null_values_to_skip = count.matching_count; + ARROW_DCHECK_EQ(count.processed_count, 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->ConsumeBufferedValues(batch_size); + values_to_skip -= batch_size; + } + } + return num_values_to_skip - values_to_skip; + } }; // ---------------------------------------------------------------------- @@ -1070,7 +1108,9 @@ class TypedColumnReaderImpl : public TypedColumnReader, int64_t ReadBatch(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, T* values, int64_t* values_read) override; - int64_t Skip(int64_t num_values_to_skip) override; + int64_t Skip(int64_t num_values_to_skip) override { + return ColumnReaderImplBase::Skip(num_values_to_skip); + } Type::type type() const override { return this->descr_->physical_type(); } @@ -1231,45 +1271,6 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size, int16_t* def return total_values; } -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 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 { - // 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 = - this->definition_level_decoder_.CountUpTo(this->max_def_level(), batch_size); - non_null_values_to_skip = count.matching_count; - ARROW_DCHECK_EQ(count.processed_count, 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->ConsumeBufferedValues(batch_size); - values_to_skip -= batch_size; - } - } - return num_values_to_skip - values_to_skip; -} - } // namespace // ---------------------------------------------------------------------- From 7984fb9a79b5e7b32c80e0135ba404ad95accd4e Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 9 Jul 2026 11:43:14 +0200 Subject: [PATCH 02/33] Break TypedRecordReader inheritance on TypedColumnReader --- cpp/src/parquet/column_reader.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 84d95f05d0e..a172e4664d5 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1318,15 +1318,14 @@ namespace internal { namespace { template -class TypedRecordReader : public TypedColumnReaderImpl, +class TypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader { public: using T = typename DType::c_type; - using BASE = TypedColumnReaderImpl; + using Base = ColumnReaderImplBase; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable) - // Pager must be set using SetPageReader. - : BASE(descr, /* pager = */ nullptr, pool) { + : Base(descr, pool) { leaf_info_ = leaf_info; nullable_values_ = leaf_info_.HasNullableValues(); at_record_start_ = true; From da5ef3e0726f2102194ee8047b68718ed836ef90 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 10 Jul 2026 13:35:47 +0200 Subject: [PATCH 03/33] Add ReadValuesBuffer --- cpp/src/parquet/column_reader.cc | 275 +++++++++++++++++++------------ cpp/src/parquet/column_reader.h | 15 +- 2 files changed, 171 insertions(+), 119 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index a172e4664d5..27d9721a12f 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1317,7 +1317,128 @@ namespace internal { namespace { +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { + if (extra_size < 0) { + throw ParquetException("Negative size (corrupt file?)"); + } + int64_t target_size = -1; + if (AddWithOverflow(size, extra_size, &target_size)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (target_size >= (1LL << 62)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (capacity >= target_size) { + return capacity; + } + return bit_util::NextPower2(target_size); +} + +class ReadValuesCursor { + public: + int64_t capacity() const { return capacity_; } + + int64_t values_count() const { return values_count_; } + + void increase_values_count(int64_t extra_values) { values_count_ += extra_values; } + + int64_t fit_capacity_for_extra(int64_t extra_values) { + auto new_capacity = compute_capacity_pow2(capacity_, values_count_, extra_values); + ARROW_DCHECK_GE(new_capacity, capacity()); + return std::exchange(capacity_, new_capacity); + } + + int64_t reset_capacity() { return std::exchange(capacity_, 0); } + + private: + int64_t values_count_ = 0; + int64_t capacity_ = 0; +}; + template +class ReadValuesBuffer : private ReadValuesCursor { + public: + using value_type = typename DType::c_type; + + using ReadValuesCursor::capacity; + using ReadValuesCursor::values_count; + + explicit ReadValuesBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} + + value_type* data() const { return values_->mutable_data_as(); } + + value_type* write_start() { return data() + values_count(); } + + void mark_values_as_written(int64_t extra_values) { + increase_values_count(extra_values); + } + + std::shared_ptr ReleaseValues(MemoryPool* pool) { + // TODO should we set values_written to zero? + auto result = values_; + const auto byte_count = bytes_for_values(values_count()); + PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true)); + values_ = AllocateBuffer(pool); + reset_capacity(); + return result; + } + + void ReserveValues(int64_t extra_values) { + const auto old_capacity = fit_capacity_for_extra(extra_values); + if (capacity() > old_capacity) { + const auto byte_count = bytes_for_values(capacity()); + PARQUET_THROW_NOT_OK(values_->Resize(byte_count, /*shrink_to_fit=*/false)); + } + } + + void ResetValues() { + if (values_count() > 0) { + PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); + ReadValuesCursor::operator=({}); + } + } + + private: + std::shared_ptr<::arrow::ResizableBuffer> values_; + + static int64_t bytes_for_values(int64_t nitems) { + constexpr auto kValueByteSize = static_cast(sizeof(value_type)); + int64_t bytes = -1; + if (MultiplyWithOverflow(nitems, kValueByteSize, &bytes)) { + throw ParquetException("Total size of items too large"); + } + return bytes; + } +}; + +template +class ReadValuesNoBuffer : private ReadValuesCursor { + public: + using value_type = typename DType::c_type; + + using ReadValuesCursor::capacity; + using ReadValuesCursor::values_count; + + explicit ReadValuesNoBuffer(MemoryPool* /* pool */) {} + + value_type* data() const { return nullptr; } + + value_type* write_start() { return nullptr; } + + void mark_values_as_written(int64_t extra_values) { + increase_values_count(extra_values); + } + + std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { + return nullptr; + } + + void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } + + void ResetValues() { ReadValuesCursor::operator=({}); } +}; + +template > class TypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader { public: @@ -1325,40 +1446,24 @@ class TypedRecordReader : public ColumnReaderImplBase, using Base = ColumnReaderImplBase; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable) - : Base(descr, pool) { + : Base(descr, pool), values_(pool) { leaf_info_ = leaf_info; nullable_values_ = leaf_info_.HasNullableValues(); at_record_start_ = true; - values_written_ = 0; null_count_ = 0; - values_capacity_ = 0; levels_written_ = 0; levels_position_ = 0; levels_capacity_ = 0; read_dense_for_nullable_ = read_dense_for_nullable; - // FIXED_LEN_BYTE_ARRAY and BYTE_ARRAY values are not stored in the `values_` buffer, - // they are read directly as Arrow. - uses_values_ = (descr->physical_type() != Type::BYTE_ARRAY && - descr->physical_type() != Type::FIXED_LEN_BYTE_ARRAY); - - if (uses_values_) { - values_ = AllocateBuffer(pool); - } valid_bits_ = AllocateBuffer(pool); def_levels_ = AllocateBuffer(pool); rep_levels_ = AllocateBuffer(pool); TypedRecordReader::Reset(); } - // Compute the values capacity in bytes for the given number of elements - int64_t bytes_for_values(int64_t nitems) const { - int64_t type_size = GetTypeByteSize(this->descr_->physical_type()); - int64_t bytes_for_values = -1; - if (MultiplyWithOverflow(nitems, type_size, &bytes_for_values)) { - throw ParquetException("Total size of items too large"); - } - return bytes_for_values; - } + uint8_t* values() const final { return reinterpret_cast(values_.data()); } + + int64_t values_written() const final { return values_.values_count(); } const void* ReadDictionary(int32_t* dictionary_length) override { if (!this->current_decoder_ && !this->HasNextInternal()) { @@ -1650,22 +1755,14 @@ class TypedRecordReader : public ColumnReaderImplBase, bool has_values_to_process() const { return levels_position_ < levels_written_; } std::shared_ptr ReleaseValues() override { - if (uses_values_) { - auto result = values_; - PARQUET_THROW_NOT_OK( - result->Resize(bytes_for_values(values_written_), /*shrink_to_fit=*/true)); - values_ = AllocateBuffer(this->pool_); - values_capacity_ = 0; - return result; - } else { - return nullptr; - } + return values_.ReleaseValues(this->pool_); } std::shared_ptr ReleaseIsValid() override { if (nullable_values()) { + const auto bit_count = bit_util::BytesForBits(values_written()); auto result = valid_bits_; - PARQUET_THROW_NOT_OK(result->Resize(bit_util::BytesForBits(values_written_), + PARQUET_THROW_NOT_OK(result->Resize(bit_count, /*shrink_to_fit=*/true)); valid_bits_ = AllocateBuffer(this->pool_); return result; @@ -1740,30 +1837,13 @@ class TypedRecordReader : public ColumnReaderImplBase, void Reserve(int64_t capacity) override { ReserveLevels(capacity); - ReserveValues(capacity); - } - - int64_t UpdateCapacity(int64_t capacity, int64_t size, int64_t extra_size) { - if (extra_size < 0) { - throw ParquetException("Negative size (corrupt file?)"); - } - int64_t target_size = -1; - if (AddWithOverflow(size, extra_size, &target_size)) { - throw ParquetException("Allocation size too large (corrupt file?)"); - } - if (target_size >= (1LL << 62)) { - throw ParquetException("Allocation size too large (corrupt file?)"); - } - if (capacity >= target_size) { - return capacity; - } - return bit_util::NextPower2(target_size); + ReserveValuesAndIsValid(capacity); } void ReserveLevels(int64_t extra_levels) { if (this->max_def_level() > 0) { const int64_t new_levels_capacity = - UpdateCapacity(levels_capacity_, levels_written_, extra_levels); + compute_capacity_pow2(levels_capacity_, levels_written_, extra_levels); if (new_levels_capacity > levels_capacity_) { constexpr auto kItemSize = static_cast(sizeof(int16_t)); int64_t capacity_in_bytes = -1; @@ -1781,22 +1861,12 @@ class TypedRecordReader : public ColumnReaderImplBase, } } - virtual void ReserveValues(int64_t extra_values) { - const int64_t new_values_capacity = - UpdateCapacity(values_capacity_, values_written_, extra_values); - if (new_values_capacity > values_capacity_) { - // XXX(wesm): A hack to avoid memory allocation when reading directly - // into builder classes - if (uses_values_) { - PARQUET_THROW_NOT_OK(values_->Resize(bytes_for_values(new_values_capacity), - /*shrink_to_fit=*/false)); - } - values_capacity_ = new_values_capacity; - } + virtual void ReserveValuesAndIsValid(int64_t extra_values) { + values_.ReserveValues(extra_values); if (nullable_values() && !read_dense_for_nullable_) { - int64_t valid_bytes_new = bit_util::BytesForBits(values_capacity_); + int64_t valid_bytes_new = bit_util::BytesForBits(values_.capacity()); if (valid_bits_->size() < valid_bytes_new) { - int64_t valid_bytes_old = bit_util::BytesForBits(values_written_); + int64_t valid_bytes_old = bit_util::BytesForBits(values_written()); PARQUET_THROW_NOT_OK( valid_bits_->Resize(valid_bytes_new, /*shrink_to_fit=*/false)); @@ -1833,17 +1903,17 @@ class TypedRecordReader : public ColumnReaderImplBase, virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) { uint8_t* valid_bits = valid_bits_->mutable_data(); - const int64_t valid_bits_offset = values_written_; + const int64_t valid_bits_offset = values_written(); int64_t num_decoded = this->current_decoder_->DecodeSpaced( - ValuesHead(), static_cast(values_with_nulls), + values_.write_start(), static_cast(values_with_nulls), static_cast(null_count), valid_bits, valid_bits_offset); CheckNumberDecoded(num_decoded, values_with_nulls); } virtual void ReadValuesDense(int64_t values_to_read) { - int64_t num_decoded = - this->current_decoder_->Decode(ValuesHead(), static_cast(values_to_read)); + int64_t num_decoded = this->current_decoder_->Decode( + values_.write_start(), static_cast(values_to_read)); CheckNumberDecoded(num_decoded, values_to_read); } @@ -1925,7 +1995,7 @@ class TypedRecordReader : public ColumnReaderImplBase, ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = levels_position_ - start_levels_position; validity_io.valid_bits = valid_bits_->mutable_data(); - validity_io.valid_bits_offset = values_written_; + validity_io.valid_bits_offset = values_written(); DefLevelsToBitmap(def_levels() + start_levels_position, levels_position_ - start_levels_position, leaf_info_, &validity_io); @@ -1942,7 +2012,7 @@ class TypedRecordReader : public ColumnReaderImplBase, // Conservative upper bound const int64_t possible_num_values = std::max(num_records, levels_written_ - levels_position_); - ReserveValues(static_cast(possible_num_values)); + ReserveValuesAndIsValid(static_cast(possible_num_values)); const int64_t start_levels_position = levels_position_; @@ -1971,10 +2041,10 @@ class TypedRecordReader : public ColumnReaderImplBase, ARROW_DCHECK_GE(null_count, 0); if (read_dense_for_nullable_) { - values_written_ += values_to_read; + values_.mark_values_as_written(values_to_read); ARROW_DCHECK_EQ(null_count, 0); } else { - values_written_ += values_to_read + null_count; + values_.mark_values_as_written(values_to_read + null_count); null_count_ += null_count; } // Total values, including null spaces, if any @@ -2020,23 +2090,15 @@ class TypedRecordReader : public ColumnReaderImplBase, } void ResetValues() { - if (values_written_ > 0) { - // Resize to 0, but do not shrink to fit - if (uses_values_) { - PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); - } + if (values_written() > 0) { + values_.ResetValues(); PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); - values_written_ = 0; - values_capacity_ = 0; null_count_ = 0; } } - protected: - template - T* ValuesHead() { - return values_->mutable_data_as() + values_written_; - } + private: + ValuesBuffer values_; LevelInfo leaf_info_; }; @@ -2048,12 +2110,13 @@ class TypedRecordReader : public ColumnReaderImplBase, /// /// The `values_` buffer is used to store the temporary values for `Decode`, and it would /// be Reset after each `Decode` call. The `valid_bits_` buffer is never used. -class FLBARecordReader final : public TypedRecordReader, - virtual public BinaryRecordReader { +class FLBARecordReader final + : public TypedRecordReader>, + virtual public BinaryRecordReader { public: FLBARecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) - : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), + : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), byte_width_(descr_->type_length()), type_(::arrow::fixed_size_binary(byte_width_)), array_builder_(type_, pool) { @@ -2065,9 +2128,8 @@ class FLBARecordReader final : public TypedRecordReader, return ::arrow::ArrayVector{std::move(chunk)}; } - void ReserveValues(int64_t extra_values) override { - ARROW_DCHECK(!uses_values_); - TypedRecordReader::ReserveValues(extra_values); + void ReserveValuesAndIsValid(int64_t extra_values) override { + TypedRecordReader::ReserveValuesAndIsValid(extra_values); PARQUET_THROW_NOT_OK(array_builder_.Reserve(extra_values)); } @@ -2081,7 +2143,7 @@ class FLBARecordReader final : public TypedRecordReader, void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { int64_t num_decoded = this->current_decoder_->DecodeArrow( static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written_, &array_builder_); + valid_bits_->mutable_data(), values_written(), &array_builder_); CheckNumberDecoded(num_decoded, values_to_read - null_count); ResetValues(); } @@ -2099,14 +2161,14 @@ class FLBARecordReader final : public TypedRecordReader, /// /// The `values_` buffers are never used, and the `accumulator_` /// is used to store the values. -class ByteArrayChunkedRecordReader final : public TypedRecordReader, - virtual public BinaryRecordReader { +class ByteArrayChunkedRecordReader final + : public TypedRecordReader>, + virtual public BinaryRecordReader { public: ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable, const std::shared_ptr<::arrow::DataType>& arrow_type) - : TypedRecordReader(descr, leaf_info, pool, - read_dense_for_nullable) { + : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable) { ARROW_DCHECK_EQ(descr_->physical_type(), Type::BYTE_ARRAY); auto arrow_binary_type = arrow_type ? arrow_type->id() : ::arrow::Type::BINARY; switch (arrow_binary_type) { @@ -2145,9 +2207,8 @@ class ByteArrayChunkedRecordReader final : public TypedRecordReaderReserve(extra_values)); } @@ -2161,7 +2222,7 @@ class ByteArrayChunkedRecordReader final : public TypedRecordReadercurrent_decoder_->DecodeArrow( static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written_, &accumulator_); + valid_bits_->mutable_data(), values_written(), &accumulator_); CheckNumberDecoded(num_decoded, values_to_read - null_count); ResetValues(); } @@ -2176,12 +2237,13 @@ class ByteArrayChunkedRecordReader final : public TypedRecordReader, - virtual public DictionaryRecordReader { +class ByteArrayDictionaryRecordReader final + : public TypedRecordReader>, + virtual public DictionaryRecordReader { public: ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) - : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), + : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), builder_(pool) { this->read_dictionary_ = true; } @@ -2238,11 +2300,11 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_.get()); num_decoded = decoder->DecodeIndicesSpaced( static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written_, &builder_); + valid_bits_->mutable_data(), values_written(), &builder_); } else { num_decoded = this->current_decoder_->DecodeArrow( static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written_, &builder_); + valid_bits_->mutable_data(), values_written(), &builder_); } ARROW_DCHECK_EQ(num_decoded, values_to_read - null_count); // Flush values since they have been copied into the builder @@ -2261,10 +2323,11 @@ template <> void TypedRecordReader::DebugPrintState() {} template <> -void TypedRecordReader::DebugPrintState() {} +void TypedRecordReader>::DebugPrintState() {} template <> -void TypedRecordReader::DebugPrintState() {} +void TypedRecordReader>::DebugPrintState() {} std::shared_ptr MakeByteArrayRecordReader( const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index 07064a9cf37..b92e978fa30 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -363,14 +363,14 @@ class PARQUET_EXPORT RecordReader { /// \brief Decoded values, including nulls, if any /// FLBA and ByteArray types do not use this array and read into their own /// builders. - uint8_t* values() const { return values_->mutable_data(); } + virtual uint8_t* values() const = 0; /// \brief Number of values written, including space left for nulls if any. /// If this Reader was constructed with read_dense_for_nullable(), there is no space for /// nulls and null_count() will be 0. There is no read-ahead/buffering for values. For /// FLBA and ByteArray types this value reflects the values written with the last /// ReadRecords call since those readers will reset the values after each call. - int64_t values_written() const { return values_written_; } + virtual int64_t values_written() const = 0; /// \brief Number of definition / repetition levels (from those that have /// been decoded) that have been consumed inside the reader. @@ -404,17 +404,6 @@ class PARQUET_EXPORT RecordReader { bool at_record_start_; int64_t records_read_; - /// \brief Stores values. These values are populated based on each ReadRecords - /// call. No extra values are buffered for the next call. SkipRecords will not - /// add any value to this buffer. - std::shared_ptr<::arrow::ResizableBuffer> values_; - /// \brief False for FIXED_LEN_BYTE_ARRAY and BYTE_ARRAY, in which case we - /// don't allocate the values buffer and we directly read into builder classes. - bool uses_values_; - - /// \brief Values that we have read into 'values_' + 'null_count_'. - int64_t values_written_; - int64_t values_capacity_; int64_t null_count_; /// \brief Each bit corresponds to one element in 'values_' and specifies if it From f6423573723b22a933805d8c6e7d5c1c9f7ad18e Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 10 Jul 2026 13:49:46 +0200 Subject: [PATCH 04/33] Split ColumnReaderImplBase definition and declaration --- cpp/src/parquet/column_reader.cc | 594 +++++++++++++++++-------------- 1 file changed, 318 insertions(+), 276 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 27d9721a12f..9f99ab60836 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -734,24 +734,48 @@ class ColumnReaderImplBase { public: using T = typename DType::c_type; - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - definition_level_decoder_(descr->max_definition_level()), - repetition_level_decoder_(descr->max_repetition_level()), - pool_(pool), - current_decoder_(pool) {} + ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool); virtual ~ColumnReaderImplBase() = default; protected: + using DecoderType = TypedDecoder; + + const ColumnDescriptor* descr_; + std::unique_ptr pager_; + std::shared_ptr current_page_; + // No data set if full schema for this field has no optional or repeated elements + LevelDecoder definition_level_decoder_; + // 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 + // the number of encoded definition levels or encoded values. For + // non-repeated, required columns, this is equal to the number of encoded + // 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_ = 0; + // The number of values from the current data page that have been decoded + // into memory or skipped over. + int64_t num_decoded_values_ = 0; + ::arrow::MemoryPool* pool_; + SkippableTypedDecoder current_decoder_; + Encoding::type current_encoding_ = Encoding::UNKNOWN; + /// Flag to signal when a new dictionary has been set, for the benefit of + /// DictionaryRecordReader + bool new_dictionary_ = false; + // The exposed encoding + ExposedEncoding exposed_encoding_ = ExposedEncoding::NO_ENCODING; + // Map of encoding type to the respective decoder object. For example, a + // column chunk's data pages may include both dictionary-encoded and + // plain-encoded data. + std::unordered_map> decoders_; + // Read up to batch_size values from the current data page into the // pre-allocated memory T* // // @returns: the number of values read into the out buffer - int64_t ReadValues(int64_t batch_size, T* out) { - int64_t num_decoded = current_decoder_->Decode(out, static_cast(batch_size)); - return num_decoded; - } + int64_t ReadValues(int64_t batch_size, T* out); // Read up to batch_size values from the current data page into the // pre-allocated memory T*, leaving spaces for null entries according @@ -759,110 +783,23 @@ class ColumnReaderImplBase { // // @returns: the number of values read into the out buffer int64_t ReadValuesSpaced(int64_t batch_size, T* out, int64_t null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - return current_decoder_->DecodeSpaced(out, static_cast(batch_size), - static_cast(null_count), valid_bits, - valid_bits_offset); - } + uint8_t* valid_bits, int64_t valid_bits_offset); // Read multiple definition levels into preallocated memory // // Returns the number of decoded definition levels - int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels) { - if (max_def_level() == 0) { - return 0; - } - return definition_level_decoder_.Decode(static_cast(batch_size), levels); - } + int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels); - bool HasNextInternal() { - // Either there is no data page available yet, or the data page has been - // exhausted - if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { - if (!ReadNewPage() || num_buffered_values_ == 0) { - return false; - } - } - return true; - } + bool HasNextInternal(); // 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) { - return 0; - } - return repetition_level_decoder_.Decode(static_cast(batch_size), levels); - } + int64_t ReadRepetitionLevels(int64_t batch_size, int16_t* levels); // Advance to the next data page - bool ReadNewPage() { - // Loop until we find the next data page. - while (true) { - current_page_ = pager_->NextPage(); - if (!current_page_) { - // EOS - return false; - } + bool ReadNewPage(); - if (current_page_->type() == PageType::DICTIONARY_PAGE) { - ConfigureDictionary(static_cast(current_page_.get())); - continue; - } else if (current_page_->type() == PageType::DATA_PAGE) { - const auto* page = static_cast(current_page_.get()); - const int64_t levels_byte_size = InitializeLevelDecoders( - *page, page->repetition_level_encoding(), page->definition_level_encoding()); - InitializeDataDecoder(*page, levels_byte_size); - return true; - } else if (current_page_->type() == PageType::DATA_PAGE_V2) { - const auto* page = static_cast(current_page_.get()); - int64_t levels_byte_size = InitializeLevelDecodersV2(*page); - InitializeDataDecoder(*page, levels_byte_size); - return true; - } else { - // We don't know what this page type is. We're allowed to skip non-data - // pages. - continue; - } - } - return true; - } - - void ConfigureDictionary(const DictionaryPage* page) { - int encoding = static_cast(page->encoding()); - if (page->encoding() == Encoding::PLAIN_DICTIONARY || - page->encoding() == Encoding::PLAIN) { - encoding = static_cast(Encoding::RLE_DICTIONARY); - } - - auto it = decoders_.find(encoding); - if (it != decoders_.end()) { - throw ParquetException("Column cannot have more than one dictionary."); - } - - if (page->encoding() == Encoding::PLAIN_DICTIONARY || - page->encoding() == Encoding::PLAIN) { - auto dictionary = MakeTypedDecoder(Encoding::PLAIN, descr_, pool_); - dictionary->SetData(page->num_values(), page->data(), page->size()); - - // The dictionary is fully decoded during DictionaryDecoder::Init, so the - // DictionaryPage buffer is no longer required after this step - // - // TODO(wesm): investigate whether this all-or-nothing decoding of the - // dictionary makes sense and whether performance can be improved - - std::unique_ptr> decoder = MakeDictDecoder(descr_, pool_); - decoder->SetDict(dictionary.get()); - decoders_[encoding] = - std::unique_ptr(dynamic_cast(decoder.release())); - } else { - ParquetException::NYI("only plain dictionary encoding has been implemented"); - } - - new_dictionary_ = true; - current_decoder_.SetDecoder(decoders_[encoding].get()); - ARROW_DCHECK(current_decoder_); - } + void ConfigureDictionary(const DictionaryPage* page); // Initialize repetition and definition level decoders on the next data page. @@ -871,222 +808,327 @@ class ColumnReaderImplBase { // The return value helps determine the number of bytes in the encoded data. int64_t InitializeLevelDecoders(const DataPage& page, Encoding::type repetition_level_encoding, - Encoding::type definition_level_encoding) { - // Read a data page. - num_buffered_values_ = page.num_values(); - - // Have not decoded any values from the data page yet - num_decoded_values_ = 0; + Encoding::type definition_level_encoding); - const uint8_t* buffer = page.data(); - int32_t levels_byte_size = 0; - int32_t max_size = page.size(); + int64_t InitializeLevelDecodersV2(const DataPageV2& page); - // 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) { - int32_t rep_levels_bytes = repetition_level_decoder_.SetData( - repetition_level_encoding, max_rep_level(), - static_cast(num_buffered_values_), buffer, max_size); - buffer += rep_levels_bytes; - levels_byte_size += rep_levels_bytes; - max_size -= rep_levels_bytes; - } - // TODO figure a way to set max_def_level_ to 0 - // if the initial value is invalid - - // Init definition levels - if (max_def_level() > 0) { - int32_t def_levels_bytes = definition_level_decoder_.SetData( - 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; - } + // Get a decoder object for this page or create a new decoder if this is the + // first page with this encoding. + void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size); - return levels_byte_size; + // Available values in the current data page, value includes repeated values + // and nulls. + int64_t available_values_current_page() const { + return num_buffered_values_ - num_decoded_values_; } - int64_t InitializeLevelDecodersV2(const DataPageV2& page) { - // Read a data page. - num_buffered_values_ = page.num_values(); + int16_t max_def_level() const; - // Have not decoded any values from the data page yet - num_decoded_values_ = 0; - const uint8_t* buffer = page.data(); + int16_t max_rep_level() const; - const int64_t total_levels_length = - static_cast(page.repetition_levels_byte_length()) + - page.definition_levels_byte_length(); + void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } - if (total_levels_length > page.size()) { - throw ParquetException("Data page too small for levels (corrupt header?)"); - } + int64_t Skip(int64_t num_values_to_skip); +}; - if (max_rep_level() > 0) { - repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(), - max_rep_level(), - static_cast(num_buffered_values_), buffer); - } - // ARROW-17453: Even if max_rep_level_ is 0, there may still be - // repetition level bytes written and/or reported in the header by - // some writers (e.g. Athena) - buffer += page.repetition_levels_byte_length(); +template +ColumnReaderImplBase::ColumnReaderImplBase(const ColumnDescriptor* descr, + ::arrow::MemoryPool* pool) + : descr_(descr), + definition_level_decoder_(descr->max_definition_level()), + repetition_level_decoder_(descr->max_repetition_level()), + pool_(pool), + current_decoder_(pool) {} - if (max_def_level() > 0) { - definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(), - max_def_level(), - static_cast(num_buffered_values_), buffer); - } +template +int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* out) { + int64_t num_decoded = current_decoder_->Decode(out, static_cast(batch_size)); + return num_decoded; +} - return total_levels_length; - } +template +int64_t ColumnReaderImplBase::ReadValuesSpaced(int64_t batch_size, T* out, + int64_t null_count, + uint8_t* valid_bits, + int64_t valid_bits_offset) { + return current_decoder_->DecodeSpaced(out, static_cast(batch_size), + static_cast(null_count), valid_bits, + valid_bits_offset); +} - // Get a decoder object for this page or create a new decoder if this is the - // first page with this encoding. - void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size) { - const uint8_t* buffer = page.data() + levels_byte_size; - const int64_t data_size = page.size() - levels_byte_size; +template +int64_t ColumnReaderImplBase::ReadDefinitionLevels(int64_t batch_size, + int16_t* levels) { + if (max_def_level() == 0) { + return 0; + } + return definition_level_decoder_.Decode(static_cast(batch_size), levels); +} - if (data_size < 0) { - throw ParquetException("Page smaller than size of encoded levels"); +template +bool ColumnReaderImplBase::HasNextInternal() { + // Either there is no data page available yet, or the data page has been + // exhausted + if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { + if (!ReadNewPage() || num_buffered_values_ == 0) { + return false; } + } + return true; +} + +template +int64_t ColumnReaderImplBase::ReadRepetitionLevels(int64_t batch_size, + int16_t* levels) { + if (max_rep_level() == 0) { + return 0; + } + return repetition_level_decoder_.Decode(static_cast(batch_size), levels); +} - Encoding::type encoding = page.encoding(); - if (IsDictionaryIndexEncoding(encoding)) { - // Normalizing the PLAIN_DICTIONARY to RLE_DICTIONARY encoding - // in decoder. - encoding = Encoding::RLE_DICTIONARY; +template +bool ColumnReaderImplBase::ReadNewPage() { + // Loop until we find the next data page. + while (true) { + current_page_ = pager_->NextPage(); + if (!current_page_) { + // EOS + return false; } - auto it = decoders_.find(static_cast(encoding)); - if (it != decoders_.end()) { - ARROW_DCHECK(it->second.get() != nullptr); - current_decoder_.SetDecoder(it->second.get()); + if (current_page_->type() == PageType::DICTIONARY_PAGE) { + ConfigureDictionary(static_cast(current_page_.get())); + continue; + } else if (current_page_->type() == PageType::DATA_PAGE) { + const auto* page = static_cast(current_page_.get()); + const int64_t levels_byte_size = InitializeLevelDecoders( + *page, page->repetition_level_encoding(), page->definition_level_encoding()); + InitializeDataDecoder(*page, levels_byte_size); + return true; + } else if (current_page_->type() == PageType::DATA_PAGE_V2) { + const auto* page = static_cast(current_page_.get()); + int64_t levels_byte_size = InitializeLevelDecodersV2(*page); + InitializeDataDecoder(*page, levels_byte_size); + return true; } else { - switch (encoding) { - case Encoding::PLAIN: - case Encoding::BYTE_STREAM_SPLIT: - case Encoding::RLE: - case Encoding::DELTA_BINARY_PACKED: - case Encoding::DELTA_BYTE_ARRAY: - case Encoding::DELTA_LENGTH_BYTE_ARRAY: { - auto decoder = MakeTypedDecoder(encoding, descr_, pool_); - current_decoder_.SetDecoder(decoder.get()); - decoders_[static_cast(encoding)] = std::move(decoder); - break; - } - - case Encoding::RLE_DICTIONARY: - throw ParquetException("Dictionary page must be before data page."); - - default: - throw ParquetException("Unknown encoding type."); - } + // We don't know what this page type is. We're allowed to skip non-data + // pages. + continue; } - current_encoding_ = encoding; - current_decoder_->SetData(static_cast(num_buffered_values_), buffer, - static_cast(data_size)); } + return true; +} - // Available values in the current data page, value includes repeated values - // and nulls. - int64_t available_values_current_page() const { - return num_buffered_values_ - num_decoded_values_; +template +void ColumnReaderImplBase::ConfigureDictionary(const DictionaryPage* page) { + int encoding = static_cast(page->encoding()); + if (page->encoding() == Encoding::PLAIN_DICTIONARY || + page->encoding() == Encoding::PLAIN) { + encoding = static_cast(Encoding::RLE_DICTIONARY); } - int16_t max_def_level() const { - // max level indirectly part of this object storage - return definition_level_decoder_.max_level(); + auto it = decoders_.find(encoding); + if (it != decoders_.end()) { + throw ParquetException("Column cannot have more than one dictionary."); } - int16_t max_rep_level() const { - // max level indirectly part of this object storage - return repetition_level_decoder_.max_level(); + if (page->encoding() == Encoding::PLAIN_DICTIONARY || + page->encoding() == Encoding::PLAIN) { + auto dictionary = MakeTypedDecoder(Encoding::PLAIN, descr_, pool_); + dictionary->SetData(page->num_values(), page->data(), page->size()); + + // The dictionary is fully decoded during DictionaryDecoder::Init, so the + // DictionaryPage buffer is no longer required after this step + // + // TODO(wesm): investigate whether this all-or-nothing decoding of the + // dictionary makes sense and whether performance can be improved + + std::unique_ptr> decoder = MakeDictDecoder(descr_, pool_); + decoder->SetDict(dictionary.get()); + decoders_[encoding] = + std::unique_ptr(dynamic_cast(decoder.release())); + } else { + ParquetException::NYI("only plain dictionary encoding has been implemented"); } - const ColumnDescriptor* descr_; + new_dictionary_ = true; + current_decoder_.SetDecoder(decoders_[encoding].get()); + ARROW_DCHECK(current_decoder_); +} - std::unique_ptr pager_; - std::shared_ptr current_page_; +template +int64_t ColumnReaderImplBase::InitializeLevelDecoders( + const DataPage& page, Encoding::type repetition_level_encoding, + Encoding::type definition_level_encoding) { + // Read a data page. + num_buffered_values_ = page.num_values(); + + // Have not decoded any values from the data page yet + num_decoded_values_ = 0; + + const uint8_t* buffer = page.data(); + int32_t levels_byte_size = 0; + int32_t max_size = page.size(); + + // 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) { + int32_t rep_levels_bytes = repetition_level_decoder_.SetData( + repetition_level_encoding, max_rep_level(), + static_cast(num_buffered_values_), buffer, max_size); + buffer += rep_levels_bytes; + levels_byte_size += rep_levels_bytes; + max_size -= rep_levels_bytes; + } + // TODO figure a way to set max_def_level_ to 0 + // if the initial value is invalid + + // Init definition levels + if (max_def_level() > 0) { + int32_t def_levels_bytes = definition_level_decoder_.SetData( + 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; + } + + return levels_byte_size; +} - // No data set if full schema for this field has no optional or repeated elements - LevelDecoder definition_level_decoder_; +template +int64_t ColumnReaderImplBase::InitializeLevelDecodersV2(const DataPageV2& page) { + // Read a data page. + num_buffered_values_ = page.num_values(); - // No data set for flat schemas. - LevelDecoder repetition_level_decoder_; + // Have not decoded any values from the data page yet + num_decoded_values_ = 0; + const uint8_t* buffer = page.data(); - // The total number of values stored in the data page. This is the maximum of - // the number of encoded definition levels or encoded values. For - // non-repeated, required columns, this is equal to the number of encoded - // 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_ = 0; + const int64_t total_levels_length = + static_cast(page.repetition_levels_byte_length()) + + page.definition_levels_byte_length(); - // The number of values from the current data page that have been decoded - // into memory or skipped over. - int64_t num_decoded_values_ = 0; + if (total_levels_length > page.size()) { + throw ParquetException("Data page too small for levels (corrupt header?)"); + } - ::arrow::MemoryPool* pool_; + if (max_rep_level() > 0) { + repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(), + max_rep_level(), + static_cast(num_buffered_values_), buffer); + } + // ARROW-17453: Even if max_rep_level_ is 0, there may still be + // repetition level bytes written and/or reported in the header by + // some writers (e.g. Athena) + buffer += page.repetition_levels_byte_length(); - using DecoderType = TypedDecoder; - SkippableTypedDecoder current_decoder_; - Encoding::type current_encoding_ = Encoding::UNKNOWN; + if (max_def_level() > 0) { + definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(), + max_def_level(), + static_cast(num_buffered_values_), buffer); + } - /// Flag to signal when a new dictionary has been set, for the benefit of - /// DictionaryRecordReader - bool new_dictionary_ = false; + return total_levels_length; +} - // The exposed encoding - ExposedEncoding exposed_encoding_ = ExposedEncoding::NO_ENCODING; +template +void ColumnReaderImplBase::InitializeDataDecoder(const DataPage& page, + int64_t levels_byte_size) { + const uint8_t* buffer = page.data() + levels_byte_size; + const int64_t data_size = page.size() - levels_byte_size; - // Map of encoding type to the respective decoder object. For example, a - // column chunk's data pages may include both dictionary-encoded and - // plain-encoded data. - std::unordered_map> decoders_; + if (data_size < 0) { + throw ParquetException("Page smaller than size of encoded levels"); + } - void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } + Encoding::type encoding = page.encoding(); + if (IsDictionaryIndexEncoding(encoding)) { + // Normalizing the PLAIN_DICTIONARY to RLE_DICTIONARY encoding + // in decoder. + encoding = Encoding::RLE_DICTIONARY; + } - int64_t 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 && HasNextInternal()) { - // If the number of values to skip is more than the number of undecoded values, skip - // 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 { - // 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 = this->definition_level_decoder_.CountUpTo( - this->max_def_level(), batch_size); - non_null_values_to_skip = count.matching_count; - ARROW_DCHECK_EQ(count.processed_count, 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); + auto it = decoders_.find(static_cast(encoding)); + if (it != decoders_.end()) { + ARROW_DCHECK(it->second.get() != nullptr); + current_decoder_.SetDecoder(it->second.get()); + } else { + switch (encoding) { + case Encoding::PLAIN: + case Encoding::BYTE_STREAM_SPLIT: + case Encoding::RLE: + case Encoding::DELTA_BINARY_PACKED: + case Encoding::DELTA_BYTE_ARRAY: + case Encoding::DELTA_LENGTH_BYTE_ARRAY: { + auto decoder = MakeTypedDecoder(encoding, descr_, pool_); + current_decoder_.SetDecoder(decoder.get()); + decoders_[static_cast(encoding)] = std::move(decoder); + break; + } + + case Encoding::RLE_DICTIONARY: + throw ParquetException("Dictionary page must be before data page."); + + default: + throw ParquetException("Unknown encoding type."); + } + } + current_encoding_ = encoding; + current_decoder_->SetData(static_cast(num_buffered_values_), buffer, + static_cast(data_size)); +} + +template +int16_t ColumnReaderImplBase::max_def_level() const { + // max level indirectly part of this object storage + return definition_level_decoder_.max_level(); +} - this->ConsumeBufferedValues(batch_size); - values_to_skip -= batch_size; +template +int16_t ColumnReaderImplBase::max_rep_level() const { + // max level indirectly part of this object storage + return repetition_level_decoder_.max_level(); +} + +template +int64_t ColumnReaderImplBase::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 && HasNextInternal()) { + // If the number of values to skip is more than the number of undecoded values, skip + // 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 { + // 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 = + this->definition_level_decoder_.CountUpTo(this->max_def_level(), batch_size); + non_null_values_to_skip = count.matching_count; + ARROW_DCHECK_EQ(count.processed_count, 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->ConsumeBufferedValues(batch_size); + values_to_skip -= batch_size; } - return num_values_to_skip - values_to_skip; } -}; + return num_values_to_skip - values_to_skip; +} // ---------------------------------------------------------------------- // TypedColumnReader implementations From 6694d78e99e0bde5b31493482487075a964b4067 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 10 Jul 2026 16:52:58 +0200 Subject: [PATCH 05/33] Add DefRepLevelsDecoder --- cpp/src/parquet/column_reader.cc | 377 +++++++++++++++++-------------- 1 file changed, 202 insertions(+), 175 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 9f99ab60836..5d65f636ae7 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -660,6 +661,10 @@ std::unique_ptr PageReader::Open(std::shared_ptr s namespace { +/*************************** + * SkippableTypedDecoder * + ***************************/ + /// Wrapper around a `TypedDecoder` pointer to skip values. /// /// Use a scratch buffer to decode values into that buffer. @@ -726,9 +731,132 @@ class SkippableTypedDecoder { } }; -// ---------------------------------------------------------------------- -// Impl base class for TypedColumnReader and RecordReader +/************************* + * DefRepLevelsDecoder * + *************************/ + +class DefRepLevelsDecoder { + public: + DefRepLevelsDecoder(int16_t max_def, int16_t max_rep) + : def_lvl_dec_(max_def), rep_lvl_dec_(max_rep) {} + + int16_t max_def_level() const { return def_lvl_dec_.max_level(); } + int16_t max_rep_level() const { return rep_lvl_dec_.max_level(); } + + // Initialize repetition and definition level decoders on the next data page. + // If the data page includes repetition and definition levels, we + // initialize the level decoders and return the number of encoded level bytes. + // The return value helps determine the number of bytes in the encoded data. + int64_t InitializeV1(const DataPageV1& page); + int64_t InitializeV2(const DataPageV2& page); + + int ReadDefinitionLevels(int batch_size, int16_t* levels); + int ReadRepetitionLevels(int batch_size, int16_t* levels); + + /// Advance both level decoders by num_levels, without materializing them. + /// + /// @return The number of values present (non-null) among them, which is the + /// number of values to skip in the data decoder. + int AdvanceLevels(int num_levels); + + private: + LevelDecoder def_lvl_dec_; + LevelDecoder rep_lvl_dec_; +}; + +/**************************************** + * DefRepLevelsDecoder Implementation * + ****************************************/ + +inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) { + const auto num_values = static_cast(page.num_values()); + + const uint8_t* buffer = page.data(); + int32_t levels_byte_size = 0; + int32_t max_size = page.size(); + + if (max_rep_level() > 0) { + int32_t rep_levels_bytes = rep_lvl_dec_.SetData( + page.repetition_level_encoding(), max_rep_level(), num_values, buffer, max_size); + buffer += rep_levels_bytes; + levels_byte_size += rep_levels_bytes; + max_size -= rep_levels_bytes; + } + if (max_def_level() > 0) { + int32_t def_levels_bytes = def_lvl_dec_.SetData( + page.definition_level_encoding(), max_def_level(), num_values, buffer, max_size); + levels_byte_size += def_levels_bytes; + max_size -= def_levels_bytes; + } + + return levels_byte_size; +} + +inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) { + const auto num_values = static_cast(page.num_values()); + + const int64_t total_levels_length = + static_cast(page.repetition_levels_byte_length()) + + page.definition_levels_byte_length(); + if (total_levels_length > page.size()) { + throw ParquetException("Data page too small for levels (corrupt header?)"); + } + + const uint8_t* buffer = page.data(); + + if (max_rep_level() > 0) { + rep_lvl_dec_.SetDataV2(page.repetition_levels_byte_length(), max_rep_level(), + num_values, buffer); + } + // ARROW-17453: Even if max_rep_level_ is 0, there may still be + // repetition level bytes written and/or reported in the header by + // some writers (e.g. Athena) + buffer += page.repetition_levels_byte_length(); + + if (max_def_level() > 0) { + def_lvl_dec_.SetDataV2(page.definition_levels_byte_length(), max_def_level(), + num_values, buffer); + } + + return total_levels_length; +} + +int DefRepLevelsDecoder::ReadDefinitionLevels(int batch_size, int16_t* levels) { + if (max_def_level() == 0) { + return 0; + } + return def_lvl_dec_.Decode(batch_size, levels); +} + +int DefRepLevelsDecoder::ReadRepetitionLevels(int batch_size, int16_t* levels) { + if (max_rep_level() == 0) { + return 0; + } + return rep_lvl_dec_.Decode(batch_size, levels); +} + +int DefRepLevelsDecoder::AdvanceLevels(int num_levels) { + int max_count = num_levels; + // Advance the definition levels, counting how many correspond to present + // (non-null) values that must be skipped in the data decoder. + if (this->max_def_level() > 0) { + const auto count = def_lvl_dec_.CountUpTo(this->max_def_level(), num_levels); + max_count = count.matching_count; + ARROW_DCHECK_EQ(count.processed_count, num_levels); + } + // Advance the repetition levels; their values are not needed. + if (this->max_rep_level() > 0) { + rep_lvl_dec_.Skip(num_levels); + } + return max_count; +} + +/************************** + * ColumnReaderImplBase * + **************************/ + +/// Impl base class for TypedColumnReader and RecordReader template class ColumnReaderImplBase { public: @@ -741,13 +869,10 @@ class ColumnReaderImplBase { protected: using DecoderType = TypedDecoder; + DefRepLevelsDecoder levels_decoder_; const ColumnDescriptor* descr_; std::unique_ptr pager_; std::shared_ptr current_page_; - // No data set if full schema for this field has no optional or repeated elements - LevelDecoder definition_level_decoder_; - // 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 // the number of encoded definition levels or encoded values. For // non-repeated, required columns, this is equal to the number of encoded @@ -785,41 +910,25 @@ class ColumnReaderImplBase { int64_t ReadValuesSpaced(int64_t batch_size, T* out, int64_t null_count, uint8_t* valid_bits, int64_t valid_bits_offset); - // Read multiple definition levels into preallocated memory - // - // Returns the number of decoded definition levels - int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels); - bool HasNextInternal(); - // Read multiple repetition levels into preallocated memory - // Returns the number of decoded repetition levels - int64_t ReadRepetitionLevels(int64_t batch_size, int16_t* levels); - // Advance to the next data page bool ReadNewPage(); void ConfigureDictionary(const DictionaryPage* page); - // Initialize repetition and definition level decoders on the next data page. - - // If the data page includes repetition and definition levels, we - // initialize the level decoders and return the number of encoded level bytes. - // The return value helps determine the number of bytes in the encoded data. - int64_t InitializeLevelDecoders(const DataPage& page, - Encoding::type repetition_level_encoding, - Encoding::type definition_level_encoding); - - int64_t InitializeLevelDecodersV2(const DataPageV2& page); - // Get a decoder object for this page or create a new decoder if this is the // first page with this encoding. void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size); // Available values in the current data page, value includes repeated values // and nulls. - int64_t available_values_current_page() const { - return num_buffered_values_ - num_decoded_values_; + int32_t available_values_current_page() const { + // The number of values in a page fits inside an int32_t according to the spec. + const int64_t out = num_buffered_values_ - num_decoded_values_; + ARROW_DCHECK_GE(out, 0); + ARROW_DCHECK_LE(out, std::numeric_limits::max()); + return static_cast(out); } int16_t max_def_level() const; @@ -834,9 +943,8 @@ class ColumnReaderImplBase { template ColumnReaderImplBase::ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - definition_level_decoder_(descr->max_definition_level()), - repetition_level_decoder_(descr->max_repetition_level()), + : levels_decoder_(descr->max_definition_level(), descr->max_repetition_level()), + descr_(descr), pool_(pool), current_decoder_(pool) {} @@ -856,15 +964,6 @@ int64_t ColumnReaderImplBase::ReadValuesSpaced(int64_t batch_size, T* out valid_bits_offset); } -template -int64_t ColumnReaderImplBase::ReadDefinitionLevels(int64_t batch_size, - int16_t* levels) { - if (max_def_level() == 0) { - return 0; - } - return definition_level_decoder_.Decode(static_cast(batch_size), levels); -} - template bool ColumnReaderImplBase::HasNextInternal() { // Either there is no data page available yet, or the data page has been @@ -877,15 +976,6 @@ bool ColumnReaderImplBase::HasNextInternal() { return true; } -template -int64_t ColumnReaderImplBase::ReadRepetitionLevels(int64_t batch_size, - int16_t* levels) { - if (max_rep_level() == 0) { - return 0; - } - return repetition_level_decoder_.Decode(static_cast(batch_size), levels); -} - template bool ColumnReaderImplBase::ReadNewPage() { // Loop until we find the next data page. @@ -900,15 +990,18 @@ bool ColumnReaderImplBase::ReadNewPage() { ConfigureDictionary(static_cast(current_page_.get())); continue; } else if (current_page_->type() == PageType::DATA_PAGE) { - const auto* page = static_cast(current_page_.get()); - const int64_t levels_byte_size = InitializeLevelDecoders( - *page, page->repetition_level_encoding(), page->definition_level_encoding()); - InitializeDataDecoder(*page, levels_byte_size); + const auto& page = static_cast(*current_page_); + const int64_t levels_byte_size = levels_decoder_.InitializeV1(page); + num_buffered_values_ = page.num_values(); + num_decoded_values_ = 0; + InitializeDataDecoder(page, levels_byte_size); return true; } else if (current_page_->type() == PageType::DATA_PAGE_V2) { - const auto* page = static_cast(current_page_.get()); - int64_t levels_byte_size = InitializeLevelDecodersV2(*page); - InitializeDataDecoder(*page, levels_byte_size); + const auto& page = static_cast(*current_page_); + const int64_t levels_byte_size = levels_decoder_.InitializeV2(page); + num_buffered_values_ = page.num_values(); + num_decoded_values_ = 0; + InitializeDataDecoder(page, levels_byte_size); return true; } else { // We don't know what this page type is. We're allowed to skip non-data @@ -956,82 +1049,6 @@ void ColumnReaderImplBase::ConfigureDictionary(const DictionaryPage* page ARROW_DCHECK(current_decoder_); } -template -int64_t ColumnReaderImplBase::InitializeLevelDecoders( - const DataPage& page, Encoding::type repetition_level_encoding, - Encoding::type definition_level_encoding) { - // Read a data page. - num_buffered_values_ = page.num_values(); - - // Have not decoded any values from the data page yet - num_decoded_values_ = 0; - - const uint8_t* buffer = page.data(); - int32_t levels_byte_size = 0; - int32_t max_size = page.size(); - - // 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) { - int32_t rep_levels_bytes = repetition_level_decoder_.SetData( - repetition_level_encoding, max_rep_level(), - static_cast(num_buffered_values_), buffer, max_size); - buffer += rep_levels_bytes; - levels_byte_size += rep_levels_bytes; - max_size -= rep_levels_bytes; - } - // TODO figure a way to set max_def_level_ to 0 - // if the initial value is invalid - - // Init definition levels - if (max_def_level() > 0) { - int32_t def_levels_bytes = definition_level_decoder_.SetData( - 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; - } - - return levels_byte_size; -} - -template -int64_t ColumnReaderImplBase::InitializeLevelDecodersV2(const DataPageV2& page) { - // Read a data page. - num_buffered_values_ = page.num_values(); - - // Have not decoded any values from the data page yet - num_decoded_values_ = 0; - const uint8_t* buffer = page.data(); - - const int64_t total_levels_length = - static_cast(page.repetition_levels_byte_length()) + - page.definition_levels_byte_length(); - - if (total_levels_length > page.size()) { - throw ParquetException("Data page too small for levels (corrupt header?)"); - } - - if (max_rep_level() > 0) { - repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(), - max_rep_level(), - static_cast(num_buffered_values_), buffer); - } - // ARROW-17453: Even if max_rep_level_ is 0, there may still be - // repetition level bytes written and/or reported in the header by - // some writers (e.g. Athena) - buffer += page.repetition_levels_byte_length(); - - if (max_def_level() > 0) { - definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(), - max_def_level(), - static_cast(num_buffered_values_), buffer); - } - - return total_levels_length; -} - template void ColumnReaderImplBase::InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size) { @@ -1081,14 +1098,12 @@ void ColumnReaderImplBase::InitializeDataDecoder(const DataPage& page, template int16_t ColumnReaderImplBase::max_def_level() const { - // max level indirectly part of this object storage - return definition_level_decoder_.max_level(); + return levels_decoder_.max_def_level(); } template int16_t ColumnReaderImplBase::max_rep_level() const { - // max level indirectly part of this object storage - return repetition_level_decoder_.max_level(); + return levels_decoder_.max_rep_level(); } template @@ -1107,21 +1122,9 @@ int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { // 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 = - this->definition_level_decoder_.CountUpTo(this->max_def_level(), batch_size); - non_null_values_to_skip = count.matching_count; - ARROW_DCHECK_EQ(count.processed_count, batch_size); - } - // Advance the repetition levels; their values are not needed. - if (this->max_rep_level() > 0) { - this->repetition_level_decoder_.Skip(batch_size); - } + int64_t non_null = levels_decoder_.AdvanceLevels(batch_size); // Skip the corresponding data values. - this->current_decoder_.Skip(non_null_values_to_skip); + this->current_decoder_.Skip(non_null); this->ConsumeBufferedValues(batch_size); values_to_skip -= batch_size; @@ -1130,6 +1133,13 @@ int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { return num_values_to_skip - values_to_skip; } +template +constexpr T clamp_to(U val) { + constexpr U kMax = std::numeric_limits::max(); + constexpr U kMin = std::numeric_limits::min(); + return static_cast(std::clamp(val, kMin, kMax)); +} + // ---------------------------------------------------------------------- // TypedColumnReader implementations @@ -1189,15 +1199,16 @@ class TypedColumnReaderImpl : public TypedColumnReader, // Read definition and repetition levels. Also return the number of definition levels // and number of values to read. This function is called before reading values. // - // ReadLevels will throw exception when any num-levels read is not equal to the number - // of the levels can be read. - void ReadLevels(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, - int64_t* num_def_levels, int64_t* non_null_values_to_read) { + // ReadLevelsInCurrentPage will throw exception when any num-levels read is not + // equal to the number of the levels can be read. + void ReadLevelsInCurrentPage(int32_t batch_size, int16_t* def_levels, + int16_t* rep_levels, int64_t* num_def_levels, + int64_t* non_null_values_to_read) { 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) { - *num_def_levels = this->ReadDefinitionLevels(batch_size, def_levels); + *num_def_levels = + this->levels_decoder_.ReadDefinitionLevels(batch_size, def_levels); if (ARROW_PREDICT_FALSE(*num_def_levels != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1215,7 +1226,8 @@ class TypedColumnReaderImpl : public TypedColumnReader, // Not present for non-repeated fields if (this->max_rep_level() > 0 && rep_levels != nullptr) { - int64_t num_rep_levels = this->ReadRepetitionLevels(batch_size, rep_levels); + int64_t num_rep_levels = + this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels); if (batch_size != num_rep_levels) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1225,8 +1237,11 @@ class TypedColumnReaderImpl : public TypedColumnReader, template int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( - int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, int32_t* indices, + int64_t batch_size_64, int16_t* def_levels, int16_t* rep_levels, int32_t* indices, int64_t* indices_read, const T** dict, int32_t* dict_len) { + // This is only reading in the current page so this fits in an int32. + const auto batch_size = clamp_to(batch_size_64); + bool has_dict_output = dict != nullptr && dict_len != nullptr; // Similar logic as ReadValues to get pages. if (!HasNext()) { @@ -1254,7 +1269,8 @@ int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( // Similar logic as ReadValues to get def levels and rep levels. int64_t num_def_levels = 0; int64_t indices_to_read = 0; - ReadLevels(batch_size, def_levels, rep_levels, &num_def_levels, &indices_to_read); + ReadLevelsInCurrentPage(batch_size, def_levels, rep_levels, &num_def_levels, + &indices_to_read); // Read dictionary indices. *indices_read = ReadDictionaryIndices(indices_to_read, indices); @@ -1272,9 +1288,9 @@ int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( } template -int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size, int16_t* def_levels, - int16_t* rep_levels, T* values, - int64_t* values_read) { +int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size_64, + int16_t* def_levels, int16_t* rep_levels, + T* values, int64_t* values_read) { // HasNext might invoke ReadNewPage until a data page with // `available_values_current_page() > 0` is found. if (!HasNext()) { @@ -1282,15 +1298,18 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size, int16_t* def return 0; } + // This is only reading in the current page so this fits in an int32. + const auto batch_size = clamp_to(batch_size_64); + // TODO(wesm): keep reading data pages until batch_size is reached, or the // row group is finished int64_t num_def_levels = 0; // Number of non-null values to read within `num_def_levels`. int64_t non_null_values_to_read = 0; - ReadLevels(batch_size, def_levels, rep_levels, &num_def_levels, - &non_null_values_to_read); + ReadLevelsInCurrentPage(batch_size, def_levels, rep_levels, &num_def_levels, + &non_null_values_to_read); // Should not return more values than available in the current data page, - // since currently, ReadLevels would only consume level from current + // since currently, ReadLevelsInCurrentPage would only consume levels from current // data page. if (ARROW_PREDICT_FALSE(num_def_levels > this->available_values_current_page())) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1556,8 +1575,10 @@ class TypedRecordReader : public ColumnReaderImplBase, /// We perform multiple batch reads until we either exhaust the row group /// or observe the desired number of records - int64_t batch_size = - std::min(level_batch_size, this->available_values_current_page()); + const int64_t batch_size_64 = + std::min(level_batch_size, this->available_values_current_page()); + // available_values_current_page fits in int32_t + const auto batch_size = static_cast(batch_size_64); // No more data in column if (batch_size == 0) { @@ -1570,12 +1591,13 @@ class TypedRecordReader : public ColumnReaderImplBase, int16_t* def_levels = this->def_levels() + levels_written_; int16_t* rep_levels = this->rep_levels() + levels_written_; - if (ARROW_PREDICT_FALSE(this->ReadDefinitionLevels(batch_size, def_levels) != - batch_size)) { + if (ARROW_PREDICT_FALSE(this->levels_decoder_.ReadDefinitionLevels( + batch_size, def_levels) != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } if (this->max_rep_level() > 0) { - int64_t rep_levels_read = this->ReadRepetitionLevels(batch_size, rep_levels); + int64_t rep_levels_read = + this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels); if (ARROW_PREDICT_FALSE(rep_levels_read != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1585,7 +1607,7 @@ class TypedRecordReader : public ColumnReaderImplBase, records_read += ReadRecordData(num_records - records_read); } else { // No repetition and definition levels, we can read values directly - batch_size = std::min(num_records - records_read, batch_size); + const auto batch_size = std::min(num_records - records_read, batch_size_64); records_read += ReadRecordData(batch_size); } } @@ -1722,8 +1744,11 @@ class TypedRecordReader : public ColumnReaderImplBase, } // Read some more levels. - int64_t batch_size = - std::min(level_batch_size, this->available_values_current_page()); + const int64_t batch_size_64 = + std::min(level_batch_size, this->available_values_current_page()); + // available_values_current_page fits in int32_t + const auto batch_size = static_cast(batch_size_64); + // No more data in column. This must be an empty page. // If we had exhausted the last page, HasNextInternal() must have advanced // to the next page. So there must be available values to process. @@ -1738,10 +1763,12 @@ class TypedRecordReader : public ColumnReaderImplBase, int16_t* def_levels = this->def_levels() + levels_written_; int16_t* rep_levels = this->rep_levels() + levels_written_; - if (this->ReadDefinitionLevels(batch_size, def_levels) != batch_size) { + if (this->levels_decoder_.ReadDefinitionLevels(batch_size, def_levels) != + batch_size) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } - if (this->ReadRepetitionLevels(batch_size, rep_levels) != batch_size) { + if (this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels) != + batch_size) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } From d9e206e4380fc0f9ff69838f57f8c4ef3f0ed910 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 10 Jul 2026 17:20:32 +0200 Subject: [PATCH 06/33] Explicit int32_t for LevelDecoder --- cpp/src/parquet/column_reader.cc | 34 ++++++++++++++++---------------- cpp/src/parquet/column_reader.h | 31 ++++++++++++++++++----------- 2 files changed, 37 insertions(+), 28 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 5d65f636ae7..c4318d7bd0a 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -107,15 +107,15 @@ struct LevelDecoder::Impl { std::variant decoder = {}; - [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { + [[nodiscard]] int32_t GetBatch(int16_t* out, int32_t batch_size) { return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); } - [[nodiscard]] int Advance(int batch_size) { + [[nodiscard]] int32_t Advance(int32_t batch_size) { return std::visit([&](auto& dec) { return dec.Advance(batch_size); }, decoder); } - auto CountUpTo(int16_t value, int batch_size) { + auto CountUpTo(int16_t value, int32_t batch_size) { return std::visit([&](auto& dec) { return dec.CountUpTo(value, batch_size); }, decoder); } @@ -126,12 +126,12 @@ LevelDecoder::LevelDecoder(int16_t max_level) LevelDecoder::~LevelDecoder() = default; -int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, - int num_buffered_values, const uint8_t* data, - int32_t data_size) { +int32_t LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, + int32_t num_buffered_values, const uint8_t* data, + int32_t data_size) { max_level_ = max_level; num_values_remaining_ = num_buffered_values; - const int value_bit_width = bit_util::Log2(max_level + 1); + const int32_t value_bit_width = bit_util::Log2(max_level + 1); switch (encoding) { case Encoding::RLE: { @@ -149,7 +149,7 @@ int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, return 4 + num_bytes; } case Encoding::BIT_PACKED: { - int num_bits = 0; + int32_t num_bits = 0; if (MultiplyWithOverflow(num_buffered_values, value_bit_width, &num_bits)) { throw ParquetException( "Number of buffered values too large (corrupt data page?)"); @@ -173,7 +173,7 @@ int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, } void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, - int num_buffered_values, const uint8_t* data) { + int32_t num_buffered_values, const uint8_t* data) { max_level_ = max_level; // Repetition and definition levels always uses RLE encoding // in the DataPageV2 format. @@ -188,9 +188,9 @@ void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, /* value_bit_width= */ bit_util::Log2(max_level + 1)); } -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); +int32_t LevelDecoder::Decode(int32_t batch_size, int16_t* levels) { + const int32_t num_values = std::min(num_values_remaining_, batch_size); + const int32_t 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_)) { @@ -204,16 +204,16 @@ 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 = impl_->Advance(num_values); +int32_t LevelDecoder::Skip(int32_t batch_size) { + const int32_t num_values = std::min(num_values_remaining_, batch_size); + const int32_t num_advanced = impl_->Advance(num_values); ARROW_DCHECK_EQ(num_values, num_advanced); num_values_remaining_ -= num_advanced; return num_advanced; } -auto LevelDecoder::CountUpTo(int16_t value, int batch_size) -> CountUpToResult { - const int num_values = std::min(num_values_remaining_, batch_size); +auto LevelDecoder::CountUpTo(int16_t value, int32_t batch_size) -> CountUpToResult { + const int32_t num_values = std::min(num_values_remaining_, batch_size); const auto result = impl_->CountUpTo(value, num_values); ARROW_DCHECK_EQ(num_values, result.processed_count); num_values_remaining_ -= result.processed_count; diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index b92e978fa30..a9ebed5d62b 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -63,6 +63,15 @@ struct PARQUET_EXPORT DataPageStats { std::optional num_rows; }; +/// Decoder for repetition of definition level. +/// +/// This decoder is used with either a deprecated bit packed (`BIT_PACKED = 4`) +/// encoding or a mixed bit packed and RLE one (`RLE = 3`). +/// Because it take as input a single buffer, `SetData` and `Decode` are typically +/// used on each of the parquet `DataPage`. +/// The number of levels is guaranteed to fit into an `int32_t` by the specification. +/// +/// @see https://research.google.com/pubs/archive/36632.pdf class PARQUET_EXPORT LevelDecoder { public: explicit LevelDecoder(int16_t max_level = 0); @@ -72,34 +81,34 @@ class PARQUET_EXPORT LevelDecoder { /// 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); + int32_t SetData(Encoding::type encoding, int16_t max_level, int32_t num_buffered_values, + const uint8_t* data, int32_t data_size); /// Initialize the LevelDecoder state with new data from a V2 page. /// /// Repetition and definition levels in V2 pages are always RLE encoded. - void SetDataV2(int32_t num_bytes, int16_t max_level, int num_buffered_values, + void SetDataV2(int32_t num_bytes, int16_t max_level, int32_t num_buffered_values, const uint8_t* data); - /// Decode a batch of levels into an array and returns the number of levels decoded. - int Decode(int batch_size, int16_t* levels); + /// Decode a batch of levels int32_to an array and returns the number of levels decoded. + int32_t Decode(int32_t batch_size, int16_t* levels); /// Advance the decoder and throw away decoder levels. - int Skip(int batch_size); + int32_t Skip(int32_t batch_size); struct CountUpToResult { - int matching_count; - int processed_count; + int32_t matching_count; + int32_t processed_count; }; /// Advance and count the number of occurrences of `value`. /// /// The count is limited to at most the next `batch_size` items. /// @return The matching value count and number of elements that were processed. - CountUpToResult CountUpTo(int16_t value, int batch_size); + CountUpToResult CountUpTo(int16_t value, int32_t batch_size); /// Return the max level used in this decoder. - int max_level() const { return max_level_; } + int32_t max_level() const { return max_level_; } private: struct Impl; @@ -107,7 +116,7 @@ class PARQUET_EXPORT LevelDecoder { std::unique_ptr impl_; /// 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; + int32_t num_values_remaining_ = 0; int16_t max_level_; }; From feb9d0dc6d488cd1ef05a9e3906070df95f076e9 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 15 Jul 2026 10:56:21 +0200 Subject: [PATCH 07/33] Change ValuesBuffer API --- cpp/src/parquet/column_reader.cc | 46 ++++++++++++++++++-------------- 1 file changed, 26 insertions(+), 20 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index c4318d7bd0a..c04aacadaba 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1395,7 +1395,7 @@ inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t ext return bit_util::NextPower2(target_size); } -class ReadValuesCursor { +class BufferCursor { public: int64_t capacity() const { return capacity_; } @@ -1416,13 +1416,13 @@ class ReadValuesCursor { int64_t capacity_ = 0; }; -template -class ReadValuesBuffer : private ReadValuesCursor { +template +class ReadValuesBuffer : private BufferCursor { public: - using value_type = typename DType::c_type; + using value_type = T; - using ReadValuesCursor::capacity; - using ReadValuesCursor::values_count; + using BufferCursor::capacity; + using BufferCursor::values_count; explicit ReadValuesBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} @@ -1455,7 +1455,7 @@ class ReadValuesBuffer : private ReadValuesCursor { void ResetValues() { if (values_count() > 0) { PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); - ReadValuesCursor::operator=({}); + BufferCursor::operator=({}); } } @@ -1472,13 +1472,13 @@ class ReadValuesBuffer : private ReadValuesCursor { } }; -template -class ReadValuesNoBuffer : private ReadValuesCursor { +template +class ReadValuesNoBuffer : private BufferCursor { public: - using value_type = typename DType::c_type; + using value_type = T; - using ReadValuesCursor::capacity; - using ReadValuesCursor::values_count; + using BufferCursor::capacity; + using BufferCursor::values_count; explicit ReadValuesNoBuffer(MemoryPool* /* pool */) {} @@ -1496,10 +1496,11 @@ class ReadValuesNoBuffer : private ReadValuesCursor { void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } - void ResetValues() { ReadValuesCursor::operator=({}); } + void ResetValues() { BufferCursor::operator=({}); } }; -template > +template > class TypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader { public: @@ -2180,7 +2181,7 @@ class TypedRecordReader : public ColumnReaderImplBase, /// The `values_` buffer is used to store the temporary values for `Decode`, and it would /// be Reset after each `Decode` call. The `valid_bits_` buffer is never used. class FLBARecordReader final - : public TypedRecordReader>, + : public TypedRecordReader>, virtual public BinaryRecordReader { public: FLBARecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, @@ -2231,7 +2232,8 @@ class FLBARecordReader final /// The `values_` buffers are never used, and the `accumulator_` /// is used to store the values. class ByteArrayChunkedRecordReader final - : public TypedRecordReader>, + : public TypedRecordReader>, virtual public BinaryRecordReader { public: ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, @@ -2307,7 +2309,8 @@ class ByteArrayChunkedRecordReader final /// If underlying column is dictionary encoded, it will call `DecodeIndices` to read, /// otherwise it will call `DecodeArrowNonNull` to read. class ByteArrayDictionaryRecordReader final - : public TypedRecordReader>, + : public TypedRecordReader>, virtual public DictionaryRecordReader { public: ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, @@ -2392,11 +2395,14 @@ template <> void TypedRecordReader::DebugPrintState() {} template <> -void TypedRecordReader>::DebugPrintState() {} +void TypedRecordReader< + ByteArrayType, + ReadValuesNoBuffer>::DebugPrintState() {} template <> -void TypedRecordReader>::DebugPrintState() {} +void TypedRecordReader>::DebugPrintState() { +} std::shared_ptr MakeByteArrayRecordReader( const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, From 1a70b9698bbf41ef9e4c244f4cd15c2100dd3213 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 15 Jul 2026 17:24:41 +0200 Subject: [PATCH 08/33] Add StackedLevelDecoder and prepare for modularity --- cpp/src/parquet/column_reader.cc | 498 ++++++++++++++++++++----------- cpp/src/parquet/column_reader.h | 6 + 2 files changed, 336 insertions(+), 168 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index c04aacadaba..d8d08c28461 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -731,14 +731,280 @@ class SkippableTypedDecoder { } }; +/****************** + * BufferCursor * + ******************/ + +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { + if (extra_size < 0) { + throw ParquetException("Negative size (corrupt file?)"); + } + int64_t target_size = -1; + if (AddWithOverflow(size, extra_size, &target_size)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (target_size >= (1LL << 62)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (capacity >= target_size) { + return capacity; + } + return bit_util::NextPower2(target_size); +} + +class BufferCursor { + public: + int64_t capacity() const { return capacity_; } + + int64_t values_count() const { return values_count_; } + + void set_values_count(int64_t vals) { values_count_ = vals; } + + int64_t fit_capacity_for_extra(int64_t extra_values) { + auto new_capacity = compute_capacity_pow2(capacity_, values_count_, extra_values); + ARROW_DCHECK_GE(new_capacity, capacity()); + return std::exchange(capacity_, new_capacity); + } + + int64_t reset_capacity() { return std::exchange(capacity_, 0); } + + private: + int64_t values_count_ = 0; + int64_t capacity_ = 0; +}; + +/********************** + * ReadValuesBuffer * + **********************/ + +template +class ReadValuesBuffer : private BufferCursor { + public: + using value_type = T; + + using BufferCursor::capacity; + using BufferCursor::values_count; + + explicit ReadValuesBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} + + value_type* data() const { return values_->mutable_data_as(); } + + value_type* write_start() { return data() + values_count(); } + + void mark_values_as_written(int64_t extra_values) { + set_values_count(values_count() + extra_values); + } + + void delete_back(int64_t count) { + ARROW_DCHECK_LE(count, values_count()); + set_values_count(values_count() - count); + } + + std::shared_ptr ReleaseValues(MemoryPool* pool) { + // TODO should we set values_written to zero? + auto result = values_; + const auto byte_count = bytes_for_values(values_count()); + PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true)); + values_ = AllocateBuffer(pool); + reset_capacity(); + return result; + } + + void ReserveValues(int64_t extra_values) { + const auto old_capacity = fit_capacity_for_extra(extra_values); + if (capacity() > old_capacity) { + const auto byte_count = bytes_for_values(capacity()); + PARQUET_THROW_NOT_OK(values_->Resize(byte_count, /*shrink_to_fit=*/false)); + } + } + + void ResetValues() { + if (values_count() > 0) { + PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); + BufferCursor::operator=({}); + } + } + + private: + std::shared_ptr<::arrow::ResizableBuffer> values_; + + static int64_t bytes_for_values(int64_t nitems) { + constexpr auto kValueByteSize = static_cast(sizeof(value_type)); + int64_t bytes = -1; + if (MultiplyWithOverflow(nitems, kValueByteSize, &bytes)) { + throw ParquetException("Total size of items too large"); + } + return bytes; + } +}; + +/************************ + * ReadValuesNoBuffer * + ************************/ + +template +class ReadValuesNoBuffer : private BufferCursor { + public: + using value_type = T; + + using BufferCursor::capacity; + using BufferCursor::values_count; + + explicit ReadValuesNoBuffer(MemoryPool* /* pool */) {} + + value_type* data() const { return nullptr; } + + value_type* write_start() { return nullptr; } + + void mark_values_as_written(int64_t extra_values) { + set_values_count(values_count() + extra_values); + } + + std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { + return nullptr; + } + + void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } + + void ResetValues() { BufferCursor::operator=({}); } +}; + +/************************* + * StackedLevelDecoder * + *************************/ + +class StackedLevelDecoder : private LevelDecoder { + public: + using Base = LevelDecoder; + using CountUpToResult = Base::CountUpToResult; + + static constexpr int32_t kMinBatchSize = kMinLevelBatchSize; + + StackedLevelDecoder(int16_t max_level, MemoryPool* pool) + : Base{max_level}, buffer_(pool) {} + + int32_t SetData(Encoding::type encoding, int16_t max_level, int32_t num_buffered_values, + const uint8_t* data, int32_t data_size); + + void SetDataV2(int32_t num_bytes, int16_t max_level, int32_t num_buffered_values, + const uint8_t* data); + + [[nodiscard]] int32_t Decode(int32_t batch_size); + + [[nodiscard]] int32_t Skip(int32_t batch_size); + + CountUpToResult CountUpTo(int16_t value, int32_t batch_size); + + int16_t* data() const { return buffer_.data(); } + + int64_t size() const { return position_; } + + private: + ReadValuesBuffer buffer_; + int64_t position_ = 0; + + int64_t total_decoded() const { return buffer_.values_count(); } + + int32_t extra_decoded() const { + const auto count = total_decoded() - size(); + ARROW_DCHECK_LT(count, kMinBatchSize); + return static_cast(count); + } + + int16_t* extra_start() { return data() + size(); } + + void AdvanceInBuffer(int32_t batch_size) { position_ += batch_size; } + + /// Left shift the extra values in the buffer by up to the batch_size. + int32_t SkipInBuffer(int32_t batch_size); +}; + +/**************************************** + * StackedLevelDecoder Implementation * + ****************************************/ + +int32_t StackedLevelDecoder::SetData(Encoding::type encoding, int16_t max_level, + int32_t num_buffered_values, const uint8_t* data, + int32_t data_size) { + buffer_.delete_back(extra_decoded()); + return Base::SetData(encoding, max_level, num_buffered_values, data, data_size); +} + +void StackedLevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, + int32_t num_buffered_values, const uint8_t* data) { + buffer_.delete_back(extra_decoded()); + return Base::SetDataV2(num_bytes, max_level, num_buffered_values, data); +} + +int32_t StackedLevelDecoder::Decode(int32_t batch_size) { + // Advance as much as possible in the already decoded buffer. + const auto buffer_batch = std::min(batch_size, extra_decoded()); + AdvanceInBuffer(buffer_batch); + + // Check what is left to decode. + const auto decode_batch = batch_size - buffer_batch; + const auto decode_possible_batch = std::min(decode_batch, Base::remaining()); + if (decode_possible_batch == 0) { + return buffer_batch; + } + + // Decode the rest plus extra greedy values. + const auto greedy_target = std::max(kMinBatchSize, decode_batch); + const auto greedy_possible = std::min(greedy_target, Base::remaining()); + buffer_.ReserveValues(/* extra= */ greedy_possible); + const auto decoded = Base::Decode(greedy_possible, buffer_.write_start()); + buffer_.mark_values_as_written(decoded); + ARROW_DCHECK_EQ(decoded, greedy_possible); + AdvanceInBuffer(decode_possible_batch); + + return buffer_batch + decode_possible_batch; +} + +int32_t StackedLevelDecoder::SkipInBuffer(int32_t batch_size) { + const auto count = std::min(batch_size, extra_decoded()); + std::copy(extra_start() + count, extra_start() + extra_decoded(), extra_start()); + buffer_.delete_back(count); + return count; +} + +int32_t StackedLevelDecoder::Skip(int32_t batch_size) { + // Left shift the extra values in the buffer by the batch_size. + const auto buffer_skipped = SkipInBuffer(batch_size); + + // Skip the rest in the decoder + const auto batch_to_skip = batch_size - buffer_skipped; + return Base::Skip(batch_to_skip) + buffer_skipped; +} + +auto StackedLevelDecoder::CountUpTo(int16_t value, + int32_t batch_size) -> CountUpToResult { + // Count in the decoded values + const auto buffer_batch = std::min(batch_size, extra_decoded()); + const auto buffer_count = static_cast( + std::count(extra_start(), extra_start() + buffer_batch, value)); + [[maybe_unused]] const auto buffer_skipped = SkipInBuffer(batch_size); + ARROW_DCHECK_EQ(buffer_skipped, buffer_batch); + + // Count the rest in the decoder. + const auto batch_to_count = batch_size - buffer_batch; + auto out = Base::CountUpTo(value, batch_to_count); + + out.matching_count += buffer_count; + out.processed_count += buffer_batch; + return out; +} + /************************* * DefRepLevelsDecoder * *************************/ +template class DefRepLevelsDecoder { public: - DefRepLevelsDecoder(int16_t max_def, int16_t max_rep) - : def_lvl_dec_(max_def), rep_lvl_dec_(max_rep) {} + using LvlDec = Decoder; + + DefRepLevelsDecoder(LvlDec def_dec, LvlDec rep_dec) + : def_lvl_dec_(std::move(def_dec)), rep_lvl_dec_(std::move(rep_dec)) {} int16_t max_def_level() const { return def_lvl_dec_.max_level(); } int16_t max_rep_level() const { return rep_lvl_dec_.max_level(); } @@ -760,15 +1026,16 @@ class DefRepLevelsDecoder { int AdvanceLevels(int num_levels); private: - LevelDecoder def_lvl_dec_; - LevelDecoder rep_lvl_dec_; + LvlDec def_lvl_dec_; + LvlDec rep_lvl_dec_; }; /**************************************** * DefRepLevelsDecoder Implementation * ****************************************/ -inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) { +template +inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) { const auto num_values = static_cast(page.num_values()); const uint8_t* buffer = page.data(); @@ -793,7 +1060,8 @@ inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) { return levels_byte_size; } -inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) { +template +inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) { const auto num_values = static_cast(page.num_values()); const int64_t total_levels_length = @@ -822,21 +1090,24 @@ inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) { return total_levels_length; } -int DefRepLevelsDecoder::ReadDefinitionLevels(int batch_size, int16_t* levels) { +template +int DefRepLevelsDecoder::ReadDefinitionLevels(int batch_size, int16_t* levels) { if (max_def_level() == 0) { return 0; } return def_lvl_dec_.Decode(batch_size, levels); } -int DefRepLevelsDecoder::ReadRepetitionLevels(int batch_size, int16_t* levels) { +template +int DefRepLevelsDecoder::ReadRepetitionLevels(int batch_size, int16_t* levels) { if (max_rep_level() == 0) { return 0; } return rep_lvl_dec_.Decode(batch_size, levels); } -int DefRepLevelsDecoder::AdvanceLevels(int num_levels) { +template +int DefRepLevelsDecoder::AdvanceLevels(int num_levels) { int max_count = num_levels; // Advance the definition levels, counting how many correspond to present // (non-null) values that must be skipped in the data decoder. @@ -857,19 +1128,24 @@ int DefRepLevelsDecoder::AdvanceLevels(int num_levels) { **************************/ /// Impl base class for TypedColumnReader and RecordReader -template +template class ColumnReaderImplBase { public: using T = typename DType::c_type; - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool); + ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, + DefRepLevelsDecoder levels_decoder) + : levels_decoder_(std::move(levels_decoder)), + descr_(descr), + pool_(pool), + current_decoder_(pool) {} virtual ~ColumnReaderImplBase() = default; protected: using DecoderType = TypedDecoder; - DefRepLevelsDecoder levels_decoder_; + DefRepLevelsDecoder levels_decoder_; const ColumnDescriptor* descr_; std::unique_ptr pager_; std::shared_ptr current_page_; @@ -940,32 +1216,24 @@ class ColumnReaderImplBase { int64_t Skip(int64_t num_values_to_skip); }; -template -ColumnReaderImplBase::ColumnReaderImplBase(const ColumnDescriptor* descr, - ::arrow::MemoryPool* pool) - : levels_decoder_(descr->max_definition_level(), descr->max_repetition_level()), - descr_(descr), - pool_(pool), - current_decoder_(pool) {} - -template -int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* out) { +template +int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* out) { int64_t num_decoded = current_decoder_->Decode(out, static_cast(batch_size)); return num_decoded; } -template -int64_t ColumnReaderImplBase::ReadValuesSpaced(int64_t batch_size, T* out, - int64_t null_count, - uint8_t* valid_bits, - int64_t valid_bits_offset) { +template +int64_t ColumnReaderImplBase::ReadValuesSpaced(int64_t batch_size, T* out, + int64_t null_count, + uint8_t* valid_bits, + int64_t valid_bits_offset) { return current_decoder_->DecodeSpaced(out, static_cast(batch_size), static_cast(null_count), valid_bits, valid_bits_offset); } -template -bool ColumnReaderImplBase::HasNextInternal() { +template +bool ColumnReaderImplBase::HasNextInternal() { // Either there is no data page available yet, or the data page has been // exhausted if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { @@ -976,8 +1244,8 @@ bool ColumnReaderImplBase::HasNextInternal() { return true; } -template -bool ColumnReaderImplBase::ReadNewPage() { +template +bool ColumnReaderImplBase::ReadNewPage() { // Loop until we find the next data page. while (true) { current_page_ = pager_->NextPage(); @@ -1012,8 +1280,9 @@ bool ColumnReaderImplBase::ReadNewPage() { return true; } -template -void ColumnReaderImplBase::ConfigureDictionary(const DictionaryPage* page) { +template +void ColumnReaderImplBase::ConfigureDictionary( + const DictionaryPage* page) { int encoding = static_cast(page->encoding()); if (page->encoding() == Encoding::PLAIN_DICTIONARY || page->encoding() == Encoding::PLAIN) { @@ -1049,9 +1318,9 @@ void ColumnReaderImplBase::ConfigureDictionary(const DictionaryPage* page ARROW_DCHECK(current_decoder_); } -template -void ColumnReaderImplBase::InitializeDataDecoder(const DataPage& page, - int64_t levels_byte_size) { +template +void ColumnReaderImplBase::InitializeDataDecoder( + const DataPage& page, int64_t levels_byte_size) { const uint8_t* buffer = page.data() + levels_byte_size; const int64_t data_size = page.size() - levels_byte_size; @@ -1096,18 +1365,18 @@ void ColumnReaderImplBase::InitializeDataDecoder(const DataPage& page, static_cast(data_size)); } -template -int16_t ColumnReaderImplBase::max_def_level() const { +template +int16_t ColumnReaderImplBase::max_def_level() const { return levels_decoder_.max_def_level(); } -template -int16_t ColumnReaderImplBase::max_rep_level() const { +template +int16_t ColumnReaderImplBase::max_rep_level() const { return levels_decoder_.max_rep_level(); } -template -int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { +template +int64_t ColumnReaderImplBase::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 && HasNextInternal()) { @@ -1145,13 +1414,18 @@ constexpr T clamp_to(U val) { template class TypedColumnReaderImpl : public TypedColumnReader, - public ColumnReaderImplBase { + public ColumnReaderImplBase { public: using T = typename DType::c_type; TypedColumnReaderImpl(const ColumnDescriptor* descr, std::unique_ptr pager, ::arrow::MemoryPool* pool) - : ColumnReaderImplBase(descr, pool) { + : ColumnReaderImplBase( + descr, pool, + { + LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level()), + }) { this->pager_ = std::move(pager); } @@ -1161,7 +1435,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, T* values, int64_t* values_read) override; int64_t Skip(int64_t num_values_to_skip) override { - return ColumnReaderImplBase::Skip(num_values_to_skip); + return ColumnReaderImplBase::Skip(num_values_to_skip); } Type::type type() const override { return this->descr_->physical_type(); } @@ -1378,137 +1652,25 @@ namespace internal { namespace { -inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { - if (extra_size < 0) { - throw ParquetException("Negative size (corrupt file?)"); - } - int64_t target_size = -1; - if (AddWithOverflow(size, extra_size, &target_size)) { - throw ParquetException("Allocation size too large (corrupt file?)"); - } - if (target_size >= (1LL << 62)) { - throw ParquetException("Allocation size too large (corrupt file?)"); - } - if (capacity >= target_size) { - return capacity; - } - return bit_util::NextPower2(target_size); -} - -class BufferCursor { - public: - int64_t capacity() const { return capacity_; } - - int64_t values_count() const { return values_count_; } - - void increase_values_count(int64_t extra_values) { values_count_ += extra_values; } - - int64_t fit_capacity_for_extra(int64_t extra_values) { - auto new_capacity = compute_capacity_pow2(capacity_, values_count_, extra_values); - ARROW_DCHECK_GE(new_capacity, capacity()); - return std::exchange(capacity_, new_capacity); - } - - int64_t reset_capacity() { return std::exchange(capacity_, 0); } - - private: - int64_t values_count_ = 0; - int64_t capacity_ = 0; -}; - -template -class ReadValuesBuffer : private BufferCursor { - public: - using value_type = T; - - using BufferCursor::capacity; - using BufferCursor::values_count; - - explicit ReadValuesBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} - - value_type* data() const { return values_->mutable_data_as(); } - - value_type* write_start() { return data() + values_count(); } - - void mark_values_as_written(int64_t extra_values) { - increase_values_count(extra_values); - } - - std::shared_ptr ReleaseValues(MemoryPool* pool) { - // TODO should we set values_written to zero? - auto result = values_; - const auto byte_count = bytes_for_values(values_count()); - PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true)); - values_ = AllocateBuffer(pool); - reset_capacity(); - return result; - } - - void ReserveValues(int64_t extra_values) { - const auto old_capacity = fit_capacity_for_extra(extra_values); - if (capacity() > old_capacity) { - const auto byte_count = bytes_for_values(capacity()); - PARQUET_THROW_NOT_OK(values_->Resize(byte_count, /*shrink_to_fit=*/false)); - } - } - - void ResetValues() { - if (values_count() > 0) { - PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); - BufferCursor::operator=({}); - } - } - - private: - std::shared_ptr<::arrow::ResizableBuffer> values_; - - static int64_t bytes_for_values(int64_t nitems) { - constexpr auto kValueByteSize = static_cast(sizeof(value_type)); - int64_t bytes = -1; - if (MultiplyWithOverflow(nitems, kValueByteSize, &bytes)) { - throw ParquetException("Total size of items too large"); - } - return bytes; - } -}; - -template -class ReadValuesNoBuffer : private BufferCursor { - public: - using value_type = T; - - using BufferCursor::capacity; - using BufferCursor::values_count; - - explicit ReadValuesNoBuffer(MemoryPool* /* pool */) {} - - value_type* data() const { return nullptr; } - - value_type* write_start() { return nullptr; } - - void mark_values_as_written(int64_t extra_values) { - increase_values_count(extra_values); - } - - std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { - return nullptr; - } - - void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } - - void ResetValues() { BufferCursor::operator=({}); } -}; +/*********************** + * TypedRecordReader * + ***********************/ template > -class TypedRecordReader : public ColumnReaderImplBase, +class TypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader { public: using T = typename DType::c_type; - using Base = ColumnReaderImplBase; + using Base = ColumnReaderImplBase; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable) - : Base(descr, pool), values_(pool) { + : Base(descr, pool, + { + LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level()), + }), + values_(pool) { leaf_info_ = leaf_info; nullable_values_ = leaf_info_.HasNullableValues(); at_record_start_ = true; diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index a9ebed5d62b..31a75d5b438 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -76,6 +76,9 @@ class PARQUET_EXPORT LevelDecoder { public: explicit LevelDecoder(int16_t max_level = 0); + LevelDecoder(LevelDecoder&&) = default; + LevelDecoder& operator=(LevelDecoder&&) = default; + ~LevelDecoder(); /// Initialize the LevelDecoder state with new data from a legacy (V1) page. @@ -110,6 +113,9 @@ class PARQUET_EXPORT LevelDecoder { /// Return the max level used in this decoder. int32_t max_level() const { return max_level_; } + /// Return the number of values left to be decoded. + int32_t remaining() const { return num_values_remaining_; } + private: struct Impl; From fb5142986148c3c85a7ad604e5a672e5f6170eb1 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 15 Jul 2026 17:57:37 +0200 Subject: [PATCH 09/33] Remove DefRepLevelsDecoder --- cpp/src/parquet/column_reader.cc | 208 +++++++++++++------------------ 1 file changed, 88 insertions(+), 120 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index d8d08c28461..9261b25fa0c 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -994,65 +994,34 @@ auto StackedLevelDecoder::CountUpTo(int16_t value, return out; } -/************************* - * DefRepLevelsDecoder * - *************************/ - -template -class DefRepLevelsDecoder { - public: - using LvlDec = Decoder; - - DefRepLevelsDecoder(LvlDec def_dec, LvlDec rep_dec) - : def_lvl_dec_(std::move(def_dec)), rep_lvl_dec_(std::move(rep_dec)) {} - - int16_t max_def_level() const { return def_lvl_dec_.max_level(); } - int16_t max_rep_level() const { return rep_lvl_dec_.max_level(); } - - // Initialize repetition and definition level decoders on the next data page. - // If the data page includes repetition and definition levels, we - // initialize the level decoders and return the number of encoded level bytes. - // The return value helps determine the number of bytes in the encoded data. - int64_t InitializeV1(const DataPageV1& page); - int64_t InitializeV2(const DataPageV2& page); - - int ReadDefinitionLevels(int batch_size, int16_t* levels); - int ReadRepetitionLevels(int batch_size, int16_t* levels); - - /// Advance both level decoders by num_levels, without materializing them. - /// - /// @return The number of values present (non-null) among them, which is the - /// number of values to skip in the data decoder. - int AdvanceLevels(int num_levels); - - private: - LvlDec def_lvl_dec_; - LvlDec rep_lvl_dec_; -}; - -/**************************************** - * DefRepLevelsDecoder Implementation * - ****************************************/ +/************************** + * ColumnReaderImplBase * + **************************/ +/// Initialize repetition and definition level decoders on the given data page. +/// +/// If the data page includes repetition and definition levels, we initialize the level +/// decoders and return the number of encoded level bytes. +/// The return value helps determine the number of bytes in the encoded data. template -inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) { +int64_t InitializeV1Levels(const DataPageV1& page, LvlDec& def_dec, LvlDec& rep_dec) { const auto num_values = static_cast(page.num_values()); const uint8_t* buffer = page.data(); int32_t levels_byte_size = 0; int32_t max_size = page.size(); - if (max_rep_level() > 0) { - int32_t rep_levels_bytes = rep_lvl_dec_.SetData( - page.repetition_level_encoding(), max_rep_level(), num_values, buffer, max_size); + if (const auto max_rep_lvl = rep_dec.max_level(); max_rep_lvl > 0) { + const int32_t rep_levels_bytes = rep_dec.SetData( + page.repetition_level_encoding(), max_rep_lvl, num_values, buffer, max_size); buffer += rep_levels_bytes; levels_byte_size += rep_levels_bytes; max_size -= rep_levels_bytes; } - if (max_def_level() > 0) { - int32_t def_levels_bytes = def_lvl_dec_.SetData( - page.definition_level_encoding(), max_def_level(), num_values, buffer, max_size); + if (const auto max_def_lvl = def_dec.max_level(); max_def_lvl > 0) { + const int32_t def_levels_bytes = def_dec.SetData( + page.definition_level_encoding(), max_def_lvl, num_values, buffer, max_size); levels_byte_size += def_levels_bytes; max_size -= def_levels_bytes; } @@ -1060,8 +1029,13 @@ inline int64_t DefRepLevelsDecoder::InitializeV1(const DataPageV1& page) return levels_byte_size; } +/// Initialize repetition and definition level decoders on the given data page. +/// +/// If the data page includes repetition and definition levels, we initialize the level +/// decoders and return the number of encoded level bytes. +/// The return value helps determine the number of bytes in the encoded data. template -inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) { +int64_t InitializeV2Levels(const DataPageV2& page, LvlDec& def_dec, LvlDec& rep_dec) { const auto num_values = static_cast(page.num_values()); const int64_t total_levels_length = @@ -1073,60 +1047,23 @@ inline int64_t DefRepLevelsDecoder::InitializeV2(const DataPageV2& page) const uint8_t* buffer = page.data(); - if (max_rep_level() > 0) { - rep_lvl_dec_.SetDataV2(page.repetition_levels_byte_length(), max_rep_level(), - num_values, buffer); + if (const auto max_rep_lvl = rep_dec.max_level(); max_rep_lvl > 0) { + rep_dec.SetDataV2(page.repetition_levels_byte_length(), max_rep_lvl, num_values, + buffer); } // ARROW-17453: Even if max_rep_level_ is 0, there may still be // repetition level bytes written and/or reported in the header by // some writers (e.g. Athena) buffer += page.repetition_levels_byte_length(); - if (max_def_level() > 0) { - def_lvl_dec_.SetDataV2(page.definition_levels_byte_length(), max_def_level(), - num_values, buffer); + if (const auto max_def_lvl = def_dec.max_level(); max_def_lvl > 0) { + def_dec.SetDataV2(page.definition_levels_byte_length(), max_def_lvl, num_values, + buffer); } return total_levels_length; } -template -int DefRepLevelsDecoder::ReadDefinitionLevels(int batch_size, int16_t* levels) { - if (max_def_level() == 0) { - return 0; - } - return def_lvl_dec_.Decode(batch_size, levels); -} - -template -int DefRepLevelsDecoder::ReadRepetitionLevels(int batch_size, int16_t* levels) { - if (max_rep_level() == 0) { - return 0; - } - return rep_lvl_dec_.Decode(batch_size, levels); -} - -template -int DefRepLevelsDecoder::AdvanceLevels(int num_levels) { - int max_count = num_levels; - // Advance the definition levels, counting how many correspond to present - // (non-null) values that must be skipped in the data decoder. - if (this->max_def_level() > 0) { - const auto count = def_lvl_dec_.CountUpTo(this->max_def_level(), num_levels); - max_count = count.matching_count; - ARROW_DCHECK_EQ(count.processed_count, num_levels); - } - // Advance the repetition levels; their values are not needed. - if (this->max_rep_level() > 0) { - rep_lvl_dec_.Skip(num_levels); - } - return max_count; -} - -/************************** - * ColumnReaderImplBase * - **************************/ - /// Impl base class for TypedColumnReader and RecordReader template class ColumnReaderImplBase { @@ -1134,8 +1071,9 @@ class ColumnReaderImplBase { using T = typename DType::c_type; ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, - DefRepLevelsDecoder levels_decoder) - : levels_decoder_(std::move(levels_decoder)), + LvlDec def_levels_decoder, LvlDec rep_levels_decoder) + : def_levels_decoder_(std::move(def_levels_decoder)), + rep_levels_decoder_(std::move(rep_levels_decoder)), descr_(descr), pool_(pool), current_decoder_(pool) {} @@ -1145,7 +1083,22 @@ class ColumnReaderImplBase { protected: using DecoderType = TypedDecoder; - DefRepLevelsDecoder levels_decoder_; + int32_t ReadDefinitionLevels(int32_t batch_size, int16_t* levels) { + if (max_def_level() == 0) { + return 0; + } + return def_levels_decoder_.Decode(batch_size, levels); + } + + int32_t ReadRepetitionLevels(int32_t batch_size, int16_t* levels) { + if (max_rep_level() == 0) { + return 0; + } + return rep_levels_decoder_.Decode(batch_size, levels); + } + + LvlDec def_levels_decoder_; + LvlDec rep_levels_decoder_; const ColumnDescriptor* descr_; std::unique_ptr pager_; std::shared_ptr current_page_; @@ -1213,6 +1166,12 @@ class ColumnReaderImplBase { void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } + /// Advance both level decoders by num_levels, without materializing them. + /// + /// @return The number of values present (non-null) among them, which is the + /// number of values to skip in the data decoder. + int32_t AdvanceLevels(int32_t num_levels); + int64_t Skip(int64_t num_values_to_skip); }; @@ -1259,14 +1218,17 @@ bool ColumnReaderImplBase::ReadNewPage() { continue; } else if (current_page_->type() == PageType::DATA_PAGE) { const auto& page = static_cast(*current_page_); - const int64_t levels_byte_size = levels_decoder_.InitializeV1(page); + const int64_t levels_byte_size = + InitializeV1Levels(page, def_levels_decoder_, rep_levels_decoder_); num_buffered_values_ = page.num_values(); num_decoded_values_ = 0; InitializeDataDecoder(page, levels_byte_size); return true; } else if (current_page_->type() == PageType::DATA_PAGE_V2) { const auto& page = static_cast(*current_page_); - const int64_t levels_byte_size = levels_decoder_.InitializeV2(page); + const int64_t levels_byte_size = + InitializeV2Levels(page, def_levels_decoder_, rep_levels_decoder_); + num_buffered_values_ = page.num_values(); num_buffered_values_ = page.num_values(); num_decoded_values_ = 0; InitializeDataDecoder(page, levels_byte_size); @@ -1367,12 +1329,29 @@ void ColumnReaderImplBase::InitializeDataDecoder( template int16_t ColumnReaderImplBase::max_def_level() const { - return levels_decoder_.max_def_level(); + return def_levels_decoder_.max_level(); } template int16_t ColumnReaderImplBase::max_rep_level() const { - return levels_decoder_.max_rep_level(); + return rep_levels_decoder_.max_level(); +} + +template +int32_t ColumnReaderImplBase::AdvanceLevels(int32_t num_levels) { + int max_count = num_levels; + // Advance the definition levels, counting how many correspond to present + // (non-null) values that must be skipped in the data decoder. + if (this->max_def_level() > 0) { + const auto count = def_levels_decoder_.CountUpTo(this->max_def_level(), num_levels); + max_count = count.matching_count; + ARROW_DCHECK_EQ(count.processed_count, num_levels); + } + // Advance the repetition levels; their values are not needed. + if (this->max_rep_level() > 0) { + rep_levels_decoder_.Skip(num_levels); + } + return max_count; } template @@ -1391,7 +1370,7 @@ int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { // whole batch fits inside this Page and no page boundary is crossed. const int batch_size = static_cast(values_to_skip); - int64_t non_null = levels_decoder_.AdvanceLevels(batch_size); + int64_t non_null = AdvanceLevels(batch_size); // Skip the corresponding data values. this->current_decoder_.Skip(non_null); @@ -1421,11 +1400,8 @@ class TypedColumnReaderImpl : public TypedColumnReader, TypedColumnReaderImpl(const ColumnDescriptor* descr, std::unique_ptr pager, ::arrow::MemoryPool* pool) : ColumnReaderImplBase( - descr, pool, - { - LevelDecoder(descr->max_definition_level()), - LevelDecoder(descr->max_repetition_level()), - }) { + descr, pool, LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level())) { this->pager_ = std::move(pager); } @@ -1481,8 +1457,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) { - *num_def_levels = - this->levels_decoder_.ReadDefinitionLevels(batch_size, def_levels); + *num_def_levels = this->ReadDefinitionLevels(batch_size, def_levels); if (ARROW_PREDICT_FALSE(*num_def_levels != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1500,8 +1475,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, // Not present for non-repeated fields if (this->max_rep_level() > 0 && rep_levels != nullptr) { - int64_t num_rep_levels = - this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels); + int64_t num_rep_levels = this->ReadRepetitionLevels(batch_size, rep_levels); if (batch_size != num_rep_levels) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1665,11 +1639,8 @@ class TypedRecordReader : public ColumnReaderImplBase, using Base = ColumnReaderImplBase; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable) - : Base(descr, pool, - { - LevelDecoder(descr->max_definition_level()), - LevelDecoder(descr->max_repetition_level()), - }), + : Base(descr, pool, LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level())), values_(pool) { leaf_info_ = leaf_info; nullable_values_ = leaf_info_.HasNullableValues(); @@ -1754,13 +1725,12 @@ class TypedRecordReader : public ColumnReaderImplBase, int16_t* def_levels = this->def_levels() + levels_written_; int16_t* rep_levels = this->rep_levels() + levels_written_; - if (ARROW_PREDICT_FALSE(this->levels_decoder_.ReadDefinitionLevels( - batch_size, def_levels) != batch_size)) { + if (ARROW_PREDICT_FALSE(this->ReadDefinitionLevels(batch_size, def_levels) != + batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } if (this->max_rep_level() > 0) { - int64_t rep_levels_read = - this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels); + int64_t rep_levels_read = this->ReadRepetitionLevels(batch_size, rep_levels); if (ARROW_PREDICT_FALSE(rep_levels_read != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } @@ -1926,12 +1896,10 @@ class TypedRecordReader : public ColumnReaderImplBase, int16_t* def_levels = this->def_levels() + levels_written_; int16_t* rep_levels = this->rep_levels() + levels_written_; - if (this->levels_decoder_.ReadDefinitionLevels(batch_size, def_levels) != - batch_size) { + if (this->ReadDefinitionLevels(batch_size, def_levels) != batch_size) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } - if (this->levels_decoder_.ReadRepetitionLevels(batch_size, rep_levels) != - batch_size) { + if (this->ReadRepetitionLevels(batch_size, rep_levels) != batch_size) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } From af9c72f7cb983607b712220e66b03d1466db4e9e Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 16 Jul 2026 10:41:27 +0200 Subject: [PATCH 10/33] Split TypedRecordReader implementation --- cpp/src/parquet/column_reader.cc | 1172 ++++++++++++++++-------------- 1 file changed, 632 insertions(+), 540 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 9261b25fa0c..9b5f57d6249 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1660,295 +1660,39 @@ class TypedRecordReader : public ColumnReaderImplBase, int64_t values_written() const final { return values_.values_count(); } - const void* ReadDictionary(int32_t* dictionary_length) override { - if (!this->current_decoder_ && !this->HasNextInternal()) { - *dictionary_length = 0; - return nullptr; - } - // Verify the current data page is dictionary encoded. The current_encoding_ should - // have been set as RLE_DICTIONARY if the page encoding is RLE_DICTIONARY or - // PLAIN_DICTIONARY. - if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { - std::stringstream ss; - ss << "Data page is not dictionary encoded. Encoding: " - << EncodingToString(this->current_encoding_); - throw ParquetException(ss.str()); - } - auto decoder = dynamic_cast*>(this->current_decoder_.get()); - const T* dictionary = nullptr; - decoder->GetDictionary(&dictionary, dictionary_length); - return reinterpret_cast(dictionary); - } - - int64_t ReadRecords(int64_t num_records) override { - if (num_records == 0) return 0; - // Delimit records, then read values at the end - int64_t records_read = 0; - - if (has_values_to_process()) { - records_read += ReadRecordData(num_records); - } - - int64_t level_batch_size = std::max(kMinLevelBatchSize, num_records); - - // If we are in the middle of a record, we continue until reaching the - // desired number of records or the end of the current record if we've found - // enough records - while (!at_record_start_ || records_read < num_records) { - // Is there more data to read in this row group? - if (!this->HasNextInternal()) { - if (!at_record_start_) { - // We ended the row group while inside a record that we haven't seen - // the end of yet. So increment the record count for the last record in - // the row group - ++records_read; - at_record_start_ = true; - } - break; - } - - /// We perform multiple batch reads until we either exhaust the row group - /// or observe the desired number of records - const int64_t batch_size_64 = - std::min(level_batch_size, this->available_values_current_page()); - // available_values_current_page fits in int32_t - const auto batch_size = static_cast(batch_size_64); - - // No more data in column - if (batch_size == 0) { - break; - } + const void* ReadDictionary(int32_t* dictionary_length) override; - if (this->max_def_level() > 0) { - ReserveLevels(batch_size); - - int16_t* def_levels = this->def_levels() + levels_written_; - int16_t* rep_levels = this->rep_levels() + levels_written_; - - if (ARROW_PREDICT_FALSE(this->ReadDefinitionLevels(batch_size, def_levels) != - batch_size)) { - throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); - } - 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); - } - } - - levels_written_ += batch_size; - records_read += ReadRecordData(num_records - records_read); - } else { - // No repetition and definition levels, we can read values directly - const auto batch_size = std::min(num_records - records_read, batch_size_64); - records_read += ReadRecordData(batch_size); - } - } - - return records_read; - } + int64_t ReadRecords(int64_t num_records) override; // Throw away levels from start_levels_position to levels_position_. // Will update levels_position_, levels_written_, and levels_capacity_ // accordingly and move the levels to left to fill in the gap. // It will resize the buffer without releasing the memory allocation. - 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_NE(def_levels_, nullptr); - - int64_t gap = levels_position_ - start_levels_position; - if (gap == 0) return; - - int64_t levels_remaining = levels_written_ - gap; - - auto left_shift = [&](::arrow::ResizableBuffer* buffer) { - auto* data = buffer->mutable_data_as(); - std::copy(data + levels_position_, data + levels_written_, - data + start_levels_position); - PARQUET_THROW_NOT_OK(buffer->Resize(levels_remaining * sizeof(int16_t), - /*shrink_to_fit=*/false)); - }; - - left_shift(def_levels_.get()); - - if (this->max_rep_level() > 0) { - ARROW_DCHECK_NE(rep_levels_, nullptr); - left_shift(rep_levels_.get()); - } - - levels_written_ -= gap; - levels_position_ -= gap; - levels_capacity_ -= gap; - } + void ThrowAwayLevels(int64_t start_levels_position); // 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); - if (!this->has_values_to_process() || num_records == 0) return 0; - - int64_t remaining_records = levels_written_ - levels_position_; - int64_t skipped_records = std::min(num_records, remaining_records); - int64_t start_levels_position = levels_position_; - // Since there is no repetition, number of levels equals number of records. - levels_position_ += skipped_records; - - // We skipped the levels by incrementing 'levels_position_'. For values - // we do not have a buffer, so we need to read them and throw them away. - // 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()); - - // 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. - // This requires shifting the levels in the buffer to left. So this will - // update levels_position_ and levels_written_. - ThrowAwayLevels(start_levels_position); - // For values, we do not have them in buffer, so we will read them and - // throw them away. - ReadAndThrowAwayValues(values_to_read); - - // Mark the levels as read in the underlying column reader. - this->ConsumeBufferedValues(skipped_records); - - return skipped_records; - } + int64_t SkipRecordsInBufferNonRepeated(int64_t num_records); // Attempts to skip num_records from the buffer. Will throw away levels // and corresponding values for the records it skipped and consumes them from the // underlying decoder. Will advance levels_position_ and update // at_record_start_. // Returns how many records were skipped. - int64_t DelimitAndSkipRecordsInBuffer(int64_t num_records) { - if (num_records == 0) return 0; - // Look at the buffered levels, delimit them based on - // (rep_level == 0), report back how many records are in there, and - // fill in how many not-null values (def_level == max_def_level_). - // DelimitRecords updates levels_position_. - int64_t start_levels_position = levels_position_; - int64_t values_seen = 0; - int64_t skipped_records = DelimitRecords(num_records, &values_seen); - ReadAndThrowAwayValues(values_seen); - // Mark those levels and values as consumed in the underlying page. - // This must be done before we throw away levels since it updates - // levels_position_ and levels_written_. - this->ConsumeBufferedValues(levels_position_ - start_levels_position); - // Updated levels_position_ and levels_written_. - ThrowAwayLevels(start_levels_position); - return skipped_records; - } + int64_t DelimitAndSkipRecordsInBuffer(int64_t num_records); // Skip records for repeated fields. For repeated fields, we are technically // reading and throwing away the levels and values since we do not know the record // boundaries in advance. Keep filling the buffer and skipping until we reach the // 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); - int64_t skipped_records = 0; - - // First consume what is in the buffer. - if (levels_position_ < levels_written_) { - // This updates at_record_start_. - skipped_records = DelimitAndSkipRecordsInBuffer(num_records); - } - - int64_t level_batch_size = - std::max(kMinLevelBatchSize, num_records - skipped_records); - - // If 'at_record_start_' is false, but (skipped_records == num_records), it - // means that for the last record that was counted, we have not seen all - // of its values yet. - while (!at_record_start_ || skipped_records < num_records) { - // Is there more data to read in this row group? - // HasNextInternal() will advance to the next page if necessary. - if (!this->HasNextInternal()) { - if (!at_record_start_) { - // We ended the row group while inside a record that we haven't seen - // the end of yet. So increment the record count for the last record - // in the row group - ++skipped_records; - at_record_start_ = true; - } - break; - } - - // Read some more levels. - const int64_t batch_size_64 = - std::min(level_batch_size, this->available_values_current_page()); - // available_values_current_page fits in int32_t - const auto batch_size = static_cast(batch_size_64); - - // No more data in column. This must be an empty page. - // If we had exhausted the last page, HasNextInternal() must have advanced - // to the next page. So there must be available values to process. - if (batch_size == 0) { - break; - } - - // For skipping we will read the levels and append them to the end - // of the def_levels and rep_levels just like for read. - ReserveLevels(batch_size); - - int16_t* def_levels = this->def_levels() + levels_written_; - int16_t* rep_levels = this->rep_levels() + levels_written_; - - if (this->ReadDefinitionLevels(batch_size, def_levels) != batch_size) { - throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); - } - if (this->ReadRepetitionLevels(batch_size, rep_levels) != batch_size) { - throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); - } - - levels_written_ += batch_size; - int64_t remaining_records = num_records - skipped_records; - // This updates at_record_start_. - skipped_records += DelimitAndSkipRecordsInBuffer(remaining_records); - } - - return skipped_records; - } + int64_t SkipRecordsRepeated(int64_t num_records); // Read 'num_values' values and throw them away. // Throws an error if it could not read 'num_values'. - void ReadAndThrowAwayValues(int64_t num_values) { - const int64_t values_read = this->current_decoder_.Skip(num_values); - if (values_read < num_values) { - std::stringstream ss; - ss << "Could not read and throw away " << num_values << " values"; - throw ParquetException(ss.str()); - } - } - - int64_t SkipRecords(int64_t num_records) override { - if (num_records == 0) return 0; + void ReadAndThrowAwayValues(int64_t num_values); - // 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) { - return this->Skip(num_records); - } - int64_t skipped_records = 0; - if (this->max_rep_level() == 0) { - // Non-repeated optional field. - // First consume whatever is in the buffer. - skipped_records = SkipRecordsInBufferNonRepeated(num_records); - - ARROW_DCHECK_LE(skipped_records, num_records); - - // For records that we have not buffered, we will use the column - // reader's Skip to do the remaining Skip. Since the field is not - // repeated number of levels to skip is the same as number of records - // to skip. - skipped_records += this->Skip(num_records - skipped_records); - } else { - skipped_records += this->SkipRecordsRepeated(num_records); - } - return skipped_records; - } + int64_t SkipRecords(int64_t num_records) override; // We may outwardly have the appearance of having exhausted a column chunk // when in fact we are in the middle of processing the last batch @@ -1958,18 +1702,7 @@ class TypedRecordReader : public ColumnReaderImplBase, return values_.ReleaseValues(this->pool_); } - std::shared_ptr ReleaseIsValid() override { - if (nullable_values()) { - const auto bit_count = bit_util::BytesForBits(values_written()); - auto result = valid_bits_; - PARQUET_THROW_NOT_OK(result->Resize(bit_count, - /*shrink_to_fit=*/true)); - valid_bits_ = AllocateBuffer(this->pool_); - return result; - } else { - return nullptr; - } - } + std::shared_ptr ReleaseIsValid() override; // Process written repetition/definition levels to reach the end of // records. Only used for repeated fields. @@ -1977,330 +1710,689 @@ class TypedRecordReader : public ColumnReaderImplBase, // number of logical records. Updates internal state of RecordReader // // \return Number of records delimited - int64_t DelimitRecords(int64_t num_records, int64_t* values_seen) { - if (ARROW_PREDICT_FALSE(num_records == 0 || levels_position_ == levels_written_)) { - *values_seen = 0; - return 0; - } - 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); - // 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 - // another record start or exhausting the ColumnChunk - int64_t level = levels_position_; - if (at_record_start_) { - if (ARROW_PREDICT_FALSE(rep_levels[levels_position_] != 0)) { - std::stringstream ss; - ss << "The repetition level at the start of a record must be 0 but got " - << rep_levels[levels_position_]; - throw ParquetException(ss.str()); - } - ++levels_position_; - // We have decided to consume the level at this position; therefore we - // must advance until we find another record boundary - at_record_start_ = false; - } + int64_t DelimitRecords(int64_t num_records, int64_t* values_seen); - // Count logical records and number of non-null values to read - ARROW_DCHECK(!at_record_start_); - // Scan repetition levels to find record end - while (levels_position_ < levels_written_) { - // We use an estimated batch size to simplify branching and - // improve performance in the common case. This might slow - // things down a bit if a single long record remains, though. - int64_t stride = - std::min(levels_written_ - levels_position_, num_records - records_read); - const int64_t position_end = levels_position_ + stride; - for (int64_t i = levels_position_; i < position_end; ++i) { - records_read += rep_levels[i] == 0; - } - levels_position_ = position_end; - if (records_read == num_records) { - // Check last rep_level reaches the boundary and - // pop the last level. - ARROW_CHECK_EQ(rep_levels[levels_position_ - 1], 0); - --levels_position_; - // We've found the number of records we were looking for. Set - // at_record_start_ to true and break + void Reserve(int64_t capacity) override; + + void ReserveLevels(int64_t extra_levels); + + virtual void ReserveValuesAndIsValid(int64_t extra_values); + + void Reset() override; + + void SetPageReader(std::unique_ptr reader) override; + + bool HasMoreData() const override { return this->pager_ != nullptr; } + + const ColumnDescriptor* descr() const override { return this->descr_; } + + // Dictionary decoders must be reset when advancing row groups + void ResetDecoders() { this->decoders_.clear(); } + + virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count); + + virtual void ReadValuesDense(int64_t values_to_read); + + // Reads repeated records and returns number of records read. Fills in + // values_to_read and null_count. + int64_t ReadRepeatedRecords(int64_t num_records, int64_t* values_to_read, + int64_t* null_count); + + // Reads optional records and returns number of records read. Fills in + // values_to_read and null_count. + int64_t ReadOptionalRecords(int64_t num_records, int64_t* values_to_read, + int64_t* null_count); + + // Reads required records and returns number of records read. Fills in + // values_to_read. + int64_t ReadRequiredRecords(int64_t num_records, int64_t* values_to_read); + + // Reads dense for optional records. First it figures out how many values to + // read. + void ReadDenseForOptional(int64_t start_levels_position, int64_t* values_to_read); + + // Reads spaced for optional or repeated fields. + void ReadSpacedForOptionalOrRepeated(int64_t start_levels_position, + int64_t* values_to_read, int64_t* null_count); + + // Return number of logical records read. + // Updates levels_position_, values_written_, and null_count_. + int64_t ReadRecordData(int64_t num_records); + + void DebugPrintState() override; + + void ResetValues(); + + private: + ValuesBuffer values_; + LevelInfo leaf_info_; +}; + +/************************************** + * TypedRecordReader Implementation * + **************************************/ + +template +const void* TypedRecordReader::ReadDictionary( + int32_t* dictionary_length) { + if (!this->current_decoder_ && !this->HasNextInternal()) { + *dictionary_length = 0; + return nullptr; + } + // Verify the current data page is dictionary encoded. The current_encoding_ should + // have been set as RLE_DICTIONARY if the page encoding is RLE_DICTIONARY or + // PLAIN_DICTIONARY. + if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { + std::stringstream ss; + ss << "Data page is not dictionary encoded. Encoding: " + << EncodingToString(this->current_encoding_); + throw ParquetException(ss.str()); + } + auto decoder = dynamic_cast*>(this->current_decoder_.get()); + const T* dictionary = nullptr; + decoder->GetDictionary(&dictionary, dictionary_length); + return reinterpret_cast(dictionary); +} + +template +int64_t TypedRecordReader::ReadRecords(int64_t num_records) { + if (num_records == 0) return 0; + // Delimit records, then read values at the end + int64_t records_read = 0; + + if (has_values_to_process()) { + records_read += ReadRecordData(num_records); + } + + int64_t level_batch_size = std::max(kMinLevelBatchSize, num_records); + + // If we are in the middle of a record, we continue until reaching the + // desired number of records or the end of the current record if we've found + // enough records + while (!at_record_start_ || records_read < num_records) { + // Is there more data to read in this row group? + if (!this->HasNextInternal()) { + if (!at_record_start_) { + // We ended the row group while inside a record that we haven't seen + // the end of yet. So increment the record count for the last record in + // the row group + ++records_read; at_record_start_ = true; - break; } + break; } - // Scan definition levels to find number of physical values - *values_seen = std::count(def_levels + level, def_levels + levels_position_, - this->max_def_level()); - return records_read; - } - void Reserve(int64_t capacity) override { - ReserveLevels(capacity); - ReserveValuesAndIsValid(capacity); - } + /// We perform multiple batch reads until we either exhaust the row group + /// or observe the desired number of records + const int64_t batch_size_64 = + std::min(level_batch_size, this->available_values_current_page()); + // available_values_current_page fits in int32_t + const auto batch_size = static_cast(batch_size_64); + + // No more data in column + if (batch_size == 0) { + break; + } - void ReserveLevels(int64_t extra_levels) { if (this->max_def_level() > 0) { - const int64_t new_levels_capacity = - compute_capacity_pow2(levels_capacity_, levels_written_, extra_levels); - if (new_levels_capacity > levels_capacity_) { - constexpr auto kItemSize = static_cast(sizeof(int16_t)); - int64_t capacity_in_bytes = -1; - if (MultiplyWithOverflow(new_levels_capacity, kItemSize, &capacity_in_bytes)) { - throw ParquetException("Allocation size too large (corrupt file?)"); - } - PARQUET_THROW_NOT_OK( - def_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); - if (this->max_rep_level() > 0) { - PARQUET_THROW_NOT_OK( - rep_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); + ReserveLevels(batch_size); + + int16_t* def_levels = this->def_levels() + levels_written_; + int16_t* rep_levels = this->rep_levels() + levels_written_; + + if (ARROW_PREDICT_FALSE(this->ReadDefinitionLevels(batch_size, def_levels) != + batch_size)) { + throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); + } + 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); } - levels_capacity_ = new_levels_capacity; } + + levels_written_ += batch_size; + records_read += ReadRecordData(num_records - records_read); + } else { + // No repetition and definition levels, we can read values directly + const auto batch_size = std::min(num_records - records_read, batch_size_64); + records_read += ReadRecordData(batch_size); } } - virtual void ReserveValuesAndIsValid(int64_t extra_values) { - values_.ReserveValues(extra_values); - if (nullable_values() && !read_dense_for_nullable_) { - int64_t valid_bytes_new = bit_util::BytesForBits(values_.capacity()); - if (valid_bits_->size() < valid_bytes_new) { - int64_t valid_bytes_old = bit_util::BytesForBits(values_written()); - PARQUET_THROW_NOT_OK( - valid_bits_->Resize(valid_bytes_new, /*shrink_to_fit=*/false)); + return records_read; +} + +template +void TypedRecordReader::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_NE(def_levels_, nullptr); + + int64_t gap = levels_position_ - start_levels_position; + if (gap == 0) return; + + int64_t levels_remaining = levels_written_ - gap; + + auto left_shift = [&](::arrow::ResizableBuffer* buffer) { + auto* data = buffer->mutable_data_as(); + std::copy(data + levels_position_, data + levels_written_, + data + start_levels_position); + PARQUET_THROW_NOT_OK(buffer->Resize(levels_remaining * sizeof(int16_t), + /*shrink_to_fit=*/false)); + }; + + left_shift(def_levels_.get()); + + if (this->max_rep_level() > 0) { + ARROW_DCHECK_NE(rep_levels_, nullptr); + left_shift(rep_levels_.get()); + } - // Avoid valgrind warnings - memset(valid_bits_->mutable_data() + valid_bytes_old, 0, - static_cast(valid_bytes_new - valid_bytes_old)); + levels_written_ -= gap; + levels_position_ -= gap; + levels_capacity_ -= gap; +} + +template +int64_t TypedRecordReader::SkipRecordsInBufferNonRepeated( + int64_t num_records) { + 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_; + int64_t skipped_records = std::min(num_records, remaining_records); + int64_t start_levels_position = levels_position_; + // Since there is no repetition, number of levels equals number of records. + levels_position_ += skipped_records; + + // We skipped the levels by incrementing 'levels_position_'. For values + // we do not have a buffer, so we need to read them and throw them away. + // 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()); + + // 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. + // This requires shifting the levels in the buffer to left. So this will + // update levels_position_ and levels_written_. + ThrowAwayLevels(start_levels_position); + // For values, we do not have them in buffer, so we will read them and + // throw them away. + ReadAndThrowAwayValues(values_to_read); + + // Mark the levels as read in the underlying column reader. + this->ConsumeBufferedValues(skipped_records); + + return skipped_records; +} + +template +int64_t TypedRecordReader::DelimitAndSkipRecordsInBuffer( + int64_t num_records) { + if (num_records == 0) return 0; + // Look at the buffered levels, delimit them based on + // (rep_level == 0), report back how many records are in there, and + // fill in how many not-null values (def_level == max_def_level_). + // DelimitRecords updates levels_position_. + int64_t start_levels_position = levels_position_; + int64_t values_seen = 0; + int64_t skipped_records = DelimitRecords(num_records, &values_seen); + ReadAndThrowAwayValues(values_seen); + // Mark those levels and values as consumed in the underlying page. + // This must be done before we throw away levels since it updates + // levels_position_ and levels_written_. + this->ConsumeBufferedValues(levels_position_ - start_levels_position); + // Updated levels_position_ and levels_written_. + ThrowAwayLevels(start_levels_position); + return skipped_records; +} + +template +int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_records) { + ARROW_DCHECK_GT(this->max_rep_level(), 0); + int64_t skipped_records = 0; + + // First consume what is in the buffer. + if (levels_position_ < levels_written_) { + // This updates at_record_start_. + skipped_records = DelimitAndSkipRecordsInBuffer(num_records); + } + + int64_t level_batch_size = + std::max(kMinLevelBatchSize, num_records - skipped_records); + + // If 'at_record_start_' is false, but (skipped_records == num_records), it + // means that for the last record that was counted, we have not seen all + // of its values yet. + while (!at_record_start_ || skipped_records < num_records) { + // Is there more data to read in this row group? + // HasNextInternal() will advance to the next page if necessary. + if (!this->HasNextInternal()) { + if (!at_record_start_) { + // We ended the row group while inside a record that we haven't seen + // the end of yet. So increment the record count for the last record + // in the row group + ++skipped_records; + at_record_start_ = true; } + break; } - } - void Reset() override { - ResetValues(); + // Read some more levels. + const int64_t batch_size_64 = + std::min(level_batch_size, this->available_values_current_page()); + // available_values_current_page fits in int32_t + const auto batch_size = static_cast(batch_size_64); - if (levels_written_ > 0) { - // Throw away levels from 0 to levels_position_. - ThrowAwayLevels(0); + // No more data in column. This must be an empty page. + // If we had exhausted the last page, HasNextInternal() must have advanced + // to the next page. So there must be available values to process. + if (batch_size == 0) { + break; } - // Call Finish on the binary builders to reset them + // For skipping we will read the levels and append them to the end + // of the def_levels and rep_levels just like for read. + ReserveLevels(batch_size); + + int16_t* def_levels = this->def_levels() + levels_written_; + int16_t* rep_levels = this->rep_levels() + levels_written_; + + if (this->ReadDefinitionLevels(batch_size, def_levels) != batch_size) { + throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); + } + if (this->ReadRepetitionLevels(batch_size, rep_levels) != batch_size) { + throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); + } + + levels_written_ += batch_size; + int64_t remaining_records = num_records - skipped_records; + // This updates at_record_start_. + skipped_records += DelimitAndSkipRecordsInBuffer(remaining_records); } - void SetPageReader(std::unique_ptr reader) override { - at_record_start_ = true; - this->pager_ = std::move(reader); - ResetDecoders(); + return skipped_records; +} + +template +void TypedRecordReader::ReadAndThrowAwayValues(int64_t num_values) { + const int64_t values_read = this->current_decoder_.Skip(num_values); + if (values_read < num_values) { + std::stringstream ss; + ss << "Could not read and throw away " << num_values << " values"; + throw ParquetException(ss.str()); } +} - bool HasMoreData() const override { return this->pager_ != nullptr; } +template +int64_t TypedRecordReader::SkipRecords(int64_t num_records) { + if (num_records == 0) return 0; - const ColumnDescriptor* descr() const override { return this->descr_; } + // 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) { + return this->Skip(num_records); + } + int64_t skipped_records = 0; + if (this->max_rep_level() == 0) { + // Non-repeated optional field. + // First consume whatever is in the buffer. + skipped_records = SkipRecordsInBufferNonRepeated(num_records); - // Dictionary decoders must be reset when advancing row groups - void ResetDecoders() { this->decoders_.clear(); } + ARROW_DCHECK_LE(skipped_records, num_records); - virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) { - uint8_t* valid_bits = valid_bits_->mutable_data(); - const int64_t valid_bits_offset = values_written(); + // For records that we have not buffered, we will use the column + // reader's Skip to do the remaining Skip. Since the field is not + // repeated number of levels to skip is the same as number of records + // to skip. + skipped_records += this->Skip(num_records - skipped_records); + } else { + skipped_records += this->SkipRecordsRepeated(num_records); + } + return skipped_records; +} - int64_t num_decoded = this->current_decoder_->DecodeSpaced( - values_.write_start(), static_cast(values_with_nulls), - static_cast(null_count), valid_bits, valid_bits_offset); - CheckNumberDecoded(num_decoded, values_with_nulls); +template +std::shared_ptr +TypedRecordReader::ReleaseIsValid() { + if (nullable_values()) { + const auto bit_count = bit_util::BytesForBits(values_written()); + auto result = valid_bits_; + PARQUET_THROW_NOT_OK(result->Resize(bit_count, + /*shrink_to_fit=*/true)); + valid_bits_ = AllocateBuffer(this->pool_); + return result; + } else { + return nullptr; } +} - virtual void ReadValuesDense(int64_t values_to_read) { - int64_t num_decoded = this->current_decoder_->Decode( - values_.write_start(), static_cast(values_to_read)); - CheckNumberDecoded(num_decoded, values_to_read); +template +int64_t TypedRecordReader::DelimitRecords(int64_t num_records, + int64_t* values_seen) { + if (ARROW_PREDICT_FALSE(num_records == 0 || levels_position_ == levels_written_)) { + *values_seen = 0; + return 0; } + 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); + // 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 + // another record start or exhausting the ColumnChunk + int64_t level = levels_position_; + if (at_record_start_) { + if (ARROW_PREDICT_FALSE(rep_levels[levels_position_] != 0)) { + std::stringstream ss; + ss << "The repetition level at the start of a record must be 0 but got " + << rep_levels[levels_position_]; + throw ParquetException(ss.str()); + } + ++levels_position_; + // We have decided to consume the level at this position; therefore we + // must advance until we find another record boundary + at_record_start_ = false; + } + + // Count logical records and number of non-null values to read + ARROW_DCHECK(!at_record_start_); + // Scan repetition levels to find record end + while (levels_position_ < levels_written_) { + // We use an estimated batch size to simplify branching and + // improve performance in the common case. This might slow + // things down a bit if a single long record remains, though. + int64_t stride = + std::min(levels_written_ - levels_position_, num_records - records_read); + const int64_t position_end = levels_position_ + stride; + for (int64_t i = levels_position_; i < position_end; ++i) { + records_read += rep_levels[i] == 0; + } + levels_position_ = position_end; + if (records_read == num_records) { + // Check last rep_level reaches the boundary and + // pop the last level. + ARROW_CHECK_EQ(rep_levels[levels_position_ - 1], 0); + --levels_position_; + // We've found the number of records we were looking for. Set + // at_record_start_ to true and break + at_record_start_ = true; + break; + } + } + // Scan definition levels to find number of physical values + *values_seen = std::count(def_levels + level, def_levels + levels_position_, + this->max_def_level()); + return records_read; +} - // Reads repeated records and returns number of records read. Fills in - // values_to_read and null_count. - int64_t ReadRepeatedRecords(int64_t num_records, int64_t* values_to_read, - int64_t* null_count) { - const int64_t start_levels_position = levels_position_; - // Note that repeated records may be required or nullable. If they have - // an optional parent in the path, they will be nullable, otherwise, - // they are required. We use leaf_info_->HasNullableValues() that looks - // at repeated_ancestor_def_level to determine if it is required or - // nullable. Even if they are required, we may have to read ahead and - // delimit the records to get the right number of values and they will - // have associated levels. - int64_t records_read = DelimitRecords(num_records, values_to_read); - if (!nullable_values() || read_dense_for_nullable_) { - ReadValuesDense(*values_to_read); - // null_count is always 0 for required. - ARROW_DCHECK_EQ(*null_count, 0); - } else { - ReadSpacedForOptionalOrRepeated(start_levels_position, values_to_read, null_count); +template +void TypedRecordReader::Reserve(int64_t capacity) { + ReserveLevels(capacity); + ReserveValuesAndIsValid(capacity); +} + +template +void TypedRecordReader::ReserveLevels(int64_t extra_levels) { + if (this->max_def_level() > 0) { + const int64_t new_levels_capacity = + compute_capacity_pow2(levels_capacity_, levels_written_, extra_levels); + if (new_levels_capacity > levels_capacity_) { + constexpr auto kItemSize = static_cast(sizeof(int16_t)); + int64_t capacity_in_bytes = -1; + if (MultiplyWithOverflow(new_levels_capacity, kItemSize, &capacity_in_bytes)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + PARQUET_THROW_NOT_OK( + def_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); + if (this->max_rep_level() > 0) { + PARQUET_THROW_NOT_OK( + rep_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); + } + levels_capacity_ = new_levels_capacity; } - return records_read; } +} - // Reads optional records and returns number of records read. Fills in - // values_to_read and null_count. - int64_t ReadOptionalRecords(int64_t num_records, int64_t* values_to_read, - int64_t* null_count) { - const int64_t start_levels_position = levels_position_; - // No repetition levels, skip delimiting logic. Each level represents a - // null or not null entry - int64_t records_read = - std::min(levels_written_ - levels_position_, num_records); - // This is advanced by DelimitRecords for the repeated field case above. - levels_position_ += records_read; - - // Optional fields are always nullable. - if (read_dense_for_nullable_) { - ReadDenseForOptional(start_levels_position, values_to_read); - // We don't need to update null_count when reading dense. It should be - // already set to 0. - ARROW_DCHECK_EQ(*null_count, 0); - } else { - ReadSpacedForOptionalOrRepeated(start_levels_position, values_to_read, null_count); +template +void TypedRecordReader::ReserveValuesAndIsValid( + int64_t extra_values) { + values_.ReserveValues(extra_values); + if (nullable_values() && !read_dense_for_nullable_) { + int64_t valid_bytes_new = bit_util::BytesForBits(values_.capacity()); + if (valid_bits_->size() < valid_bytes_new) { + int64_t valid_bytes_old = bit_util::BytesForBits(values_written()); + PARQUET_THROW_NOT_OK(valid_bits_->Resize(valid_bytes_new, /*shrink_to_fit=*/false)); + + // Avoid valgrind warnings + memset(valid_bits_->mutable_data() + valid_bytes_old, 0, + static_cast(valid_bytes_new - valid_bytes_old)); } - return records_read; } +} - // Reads required records and returns number of records read. Fills in - // values_to_read. - int64_t ReadRequiredRecords(int64_t num_records, int64_t* values_to_read) { - *values_to_read = num_records; - ReadValuesDense(*values_to_read); - return num_records; +template +void TypedRecordReader::Reset() { + ResetValues(); + + if (levels_written_ > 0) { + // Throw away levels from 0 to levels_position_. + ThrowAwayLevels(0); } - // Reads dense for optional records. First it figures out how many values to - // read. - void ReadDenseForOptional(int64_t start_levels_position, int64_t* values_to_read) { - // levels_position_ must already be incremented based on number of records - // read. - ARROW_DCHECK_GE(levels_position_, start_levels_position); - - // 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()); + // Call Finish on the binary builders to reset them +} + +template +void TypedRecordReader::SetPageReader( + std::unique_ptr reader) { + at_record_start_ = true; + this->pager_ = std::move(reader); + ResetDecoders(); +} + +template +void TypedRecordReader::ReadValuesSpaced(int64_t values_with_nulls, + int64_t null_count) { + uint8_t* valid_bits = valid_bits_->mutable_data(); + const int64_t valid_bits_offset = values_written(); + + int64_t num_decoded = this->current_decoder_->DecodeSpaced( + values_.write_start(), static_cast(values_with_nulls), + static_cast(null_count), valid_bits, valid_bits_offset); + CheckNumberDecoded(num_decoded, values_with_nulls); +} + +template +void TypedRecordReader::ReadValuesDense(int64_t values_to_read) { + int64_t num_decoded = this->current_decoder_->Decode(values_.write_start(), + static_cast(values_to_read)); + CheckNumberDecoded(num_decoded, values_to_read); +} + +template +int64_t TypedRecordReader::ReadRepeatedRecords( + int64_t num_records, int64_t* values_to_read, int64_t* null_count) { + const int64_t start_levels_position = levels_position_; + // Note that repeated records may be required or nullable. If they have + // an optional parent in the path, they will be nullable, otherwise, + // they are required. We use leaf_info_->HasNullableValues() that looks + // at repeated_ancestor_def_level to determine if it is required or + // nullable. Even if they are required, we may have to read ahead and + // delimit the records to get the right number of values and they will + // have associated levels. + int64_t records_read = DelimitRecords(num_records, values_to_read); + if (!nullable_values() || read_dense_for_nullable_) { ReadValuesDense(*values_to_read); + // null_count is always 0 for required. + ARROW_DCHECK_EQ(*null_count, 0); + } else { + ReadSpacedForOptionalOrRepeated(start_levels_position, values_to_read, null_count); } + return records_read; +} - // Reads spaced for optional or repeated fields. - void ReadSpacedForOptionalOrRepeated(int64_t start_levels_position, - int64_t* values_to_read, int64_t* null_count) { - // levels_position_ must already be incremented based on number of records - // read. - ARROW_DCHECK_GE(levels_position_, start_levels_position); - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = levels_position_ - start_levels_position; - validity_io.valid_bits = valid_bits_->mutable_data(); - validity_io.valid_bits_offset = values_written(); - - DefLevelsToBitmap(def_levels() + start_levels_position, - levels_position_ - start_levels_position, leaf_info_, &validity_io); - *values_to_read = validity_io.values_read - validity_io.null_count; - *null_count = validity_io.null_count; - ARROW_DCHECK_GE(*values_to_read, 0); - ARROW_DCHECK_GE(*null_count, 0); - ReadValuesSpaced(validity_io.values_read, *null_count); +template +int64_t TypedRecordReader::ReadOptionalRecords( + int64_t num_records, int64_t* values_to_read, int64_t* null_count) { + const int64_t start_levels_position = levels_position_; + // No repetition levels, skip delimiting logic. Each level represents a + // null or not null entry + int64_t records_read = + std::min(levels_written_ - levels_position_, num_records); + // This is advanced by DelimitRecords for the repeated field case above. + levels_position_ += records_read; + + // Optional fields are always nullable. + if (read_dense_for_nullable_) { + ReadDenseForOptional(start_levels_position, values_to_read); + // We don't need to update null_count when reading dense. It should be + // already set to 0. + ARROW_DCHECK_EQ(*null_count, 0); + } else { + ReadSpacedForOptionalOrRepeated(start_levels_position, values_to_read, null_count); } + return records_read; +} - // Return number of logical records read. - // Updates levels_position_, values_written_, and null_count_. - int64_t ReadRecordData(int64_t num_records) { - // Conservative upper bound - const int64_t possible_num_values = - std::max(num_records, levels_written_ - levels_position_); - ReserveValuesAndIsValid(static_cast(possible_num_values)); - - const int64_t start_levels_position = levels_position_; - - // To be updated by the function calls below for each of the repetition - // types. - int64_t records_read = 0; - int64_t values_to_read = 0; - int64_t null_count = 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) { - // Non-repeated optional values are always nullable. - // This call updates levels_position_. - ARROW_DCHECK(nullable_values()); - records_read = ReadOptionalRecords(num_records, &values_to_read, &null_count); - } else { - ARROW_DCHECK(!nullable_values()); - records_read = ReadRequiredRecords(num_records, &values_to_read); - // We don't need to update null_count, since it is 0. - } +template +int64_t TypedRecordReader::ReadRequiredRecords( + int64_t num_records, int64_t* values_to_read) { + *values_to_read = num_records; + ReadValuesDense(*values_to_read); + return num_records; +} - ARROW_DCHECK_GE(records_read, 0); - ARROW_DCHECK_GE(values_to_read, 0); - ARROW_DCHECK_GE(null_count, 0); +template +void TypedRecordReader::ReadDenseForOptional( + int64_t start_levels_position, int64_t* values_to_read) { + // levels_position_ must already be incremented based on number of records + // read. + ARROW_DCHECK_GE(levels_position_, start_levels_position); - if (read_dense_for_nullable_) { - values_.mark_values_as_written(values_to_read); - ARROW_DCHECK_EQ(null_count, 0); - } else { - values_.mark_values_as_written(values_to_read + null_count); - null_count_ += null_count; - } - // Total values, including null spaces, if any - if (this->max_def_level() > 0) { - // Optional, repeated, or some mix thereof - this->ConsumeBufferedValues(levels_position_ - start_levels_position); - } else { - // Flat, non-repeated - this->ConsumeBufferedValues(values_to_read); - } + // 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()); + ReadValuesDense(*values_to_read); +} - return records_read; +template +void TypedRecordReader::ReadSpacedForOptionalOrRepeated( + int64_t start_levels_position, int64_t* values_to_read, int64_t* null_count) { + // levels_position_ must already be incremented based on number of records + // read. + ARROW_DCHECK_GE(levels_position_, start_levels_position); + ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = levels_position_ - start_levels_position; + validity_io.valid_bits = valid_bits_->mutable_data(); + validity_io.valid_bits_offset = values_written(); + + DefLevelsToBitmap(def_levels() + start_levels_position, + levels_position_ - start_levels_position, leaf_info_, &validity_io); + *values_to_read = validity_io.values_read - validity_io.null_count; + *null_count = validity_io.null_count; + ARROW_DCHECK_GE(*values_to_read, 0); + ARROW_DCHECK_GE(*null_count, 0); + ReadValuesSpaced(validity_io.values_read, *null_count); +} + +template +int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { + // Conservative upper bound + const int64_t possible_num_values = + std::max(num_records, levels_written_ - levels_position_); + ReserveValuesAndIsValid(static_cast(possible_num_values)); + + const int64_t start_levels_position = levels_position_; + + // To be updated by the function calls below for each of the repetition + // types. + int64_t records_read = 0; + int64_t values_to_read = 0; + int64_t null_count = 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) { + // Non-repeated optional values are always nullable. + // This call updates levels_position_. + ARROW_DCHECK(nullable_values()); + records_read = ReadOptionalRecords(num_records, &values_to_read, &null_count); + } else { + ARROW_DCHECK(!nullable_values()); + records_read = ReadRequiredRecords(num_records, &values_to_read); + // We don't need to update null_count, since it is 0. } - void DebugPrintState() override { - const int16_t* def_levels = this->def_levels(); - const int16_t* rep_levels = this->rep_levels(); - const int64_t total_levels_read = levels_position_; + ARROW_DCHECK_GE(records_read, 0); + ARROW_DCHECK_GE(values_to_read, 0); + ARROW_DCHECK_GE(null_count, 0); - const T* vals = reinterpret_cast(this->values()); + if (read_dense_for_nullable_) { + values_.mark_values_as_written(values_to_read); + ARROW_DCHECK_EQ(null_count, 0); + } else { + values_.mark_values_as_written(values_to_read + null_count); + null_count_ += null_count; + } + // Total values, including null spaces, if any + if (this->max_def_level() > 0) { + // Optional, repeated, or some mix thereof + this->ConsumeBufferedValues(levels_position_ - start_levels_position); + } else { + // Flat, non-repeated + this->ConsumeBufferedValues(values_to_read); + } - if (leaf_info_.def_level > 0) { - std::cout << "def levels: "; - for (int64_t i = 0; i < total_levels_read; ++i) { - std::cout << def_levels[i] << " "; - } - std::cout << std::endl; - } + return records_read; +} - if (leaf_info_.rep_level > 0) { - std::cout << "rep levels: "; - for (int64_t i = 0; i < total_levels_read; ++i) { - std::cout << rep_levels[i] << " "; - } - std::cout << std::endl; - } +template +void TypedRecordReader::DebugPrintState() { + const int16_t* def_levels = this->def_levels(); + const int16_t* rep_levels = this->rep_levels(); + const int64_t total_levels_read = levels_position_; + + const T* vals = reinterpret_cast(this->values()); - std::cout << "values: "; - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << vals[i] << " "; + if (leaf_info_.def_level > 0) { + std::cout << "def levels: "; + for (int64_t i = 0; i < total_levels_read; ++i) { + std::cout << def_levels[i] << " "; } std::cout << std::endl; } - void ResetValues() { - if (values_written() > 0) { - values_.ResetValues(); - PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); - null_count_ = 0; + if (leaf_info_.rep_level > 0) { + std::cout << "rep levels: "; + for (int64_t i = 0; i < total_levels_read; ++i) { + std::cout << rep_levels[i] << " "; } + std::cout << std::endl; } - private: - ValuesBuffer values_; - LevelInfo leaf_info_; -}; + std::cout << "values: "; + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << vals[i] << " "; + } + std::cout << std::endl; +} + +template +void TypedRecordReader::ResetValues() { + if (values_written() > 0) { + values_.ResetValues(); + PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); + null_count_ = 0; + } +} /// In FLBARecordReader, we read fixed length byte array values. /// From 21340c79cf7c9d6f8cd38bbbbca88a7935301259 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 16 Jul 2026 14:03:13 +0200 Subject: [PATCH 11/33] First RequiredTypedRecordReader --- cpp/src/parquet/column_reader.cc | 191 ++++++++++++++++++++++++++++--- 1 file changed, 174 insertions(+), 17 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 9b5f57d6249..436660fc59f 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -25,20 +25,15 @@ #include #include #include -#include #include #include #include #include -#include "arrow/array.h" -#include "arrow/array/array_binary.h" #include "arrow/array/builder_binary.h" #include "arrow/array/builder_dict.h" -#include "arrow/array/builder_primitive.h" #include "arrow/chunked_array.h" #include "arrow/type.h" -#include "arrow/util/bit_stream_utils_internal.h" #include "arrow/util/bit_util.h" #include "arrow/util/checked_cast.h" #include "arrow/util/compression.h" @@ -2394,6 +2389,154 @@ void TypedRecordReader::ResetValues() { } } +/******************************* + * RequiredTypedRecordReader * + *******************************/ + +template +class RequiredTypedRecordReader final : public ColumnReaderImplBase, + public RecordReader { + public: + using T = typename DType::c_type; + using Base = ColumnReaderImplBase; + using ValuesBuffer = ReadValuesBuffer; + + RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool) + : Base(descr, pool, LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level())), + values_(pool) { + nullable_values_ = false; + at_record_start_ = true; + null_count_ = 0; + levels_written_ = 0; + levels_position_ = 0; + levels_capacity_ = 0; + read_dense_for_nullable_ = false; + valid_bits_ = AllocateBuffer(pool); + RequiredTypedRecordReader::Reset(); + } + + uint8_t* values() const final { return reinterpret_cast(values_.data()); } + + int64_t values_written() const final { return values_.values_count(); } + + const void* ReadDictionary(int32_t* dictionary_length) override; + + int64_t ReadRecords(int64_t num_records) override; + + int64_t SkipRecords(int64_t num_records) override { return this->Skip(num_records); } + + std::shared_ptr ReleaseValues() override { + return values_.ReleaseValues(this->pool_); + } + + std::shared_ptr ReleaseIsValid() override { return nullptr; } + + void Reserve(int64_t extra_values) override { values_.ReserveValues(extra_values); } + + void Reset() override; + + void SetPageReader(std::unique_ptr reader) override; + + bool HasMoreData() const override { return this->pager_ != nullptr; } + + const ColumnDescriptor* descr() const override { return this->descr_; } + + // Dictionary decoders must be reset when advancing row groups + void ResetDecoders() { this->decoders_.clear(); } + + void DebugPrintState() override; + + private: + ValuesBuffer values_; +}; + +/********************************************** + * RequiredTypedRecordReader Implementation * + **********************************************/ + +template +const void* RequiredTypedRecordReader::ReadDictionary(int32_t* dictionary_length) { + if (!this->current_decoder_ && !this->HasNextInternal()) { + *dictionary_length = 0; + return nullptr; + } + // Verify the current data page is dictionary encoded. The current_encoding_ should + // have been set as RLE_DICTIONARY if the page encoding is RLE_DICTIONARY or + // PLAIN_DICTIONARY. + if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { + std::stringstream ss; + ss << "Data page is not dictionary encoded. Encoding: " + << EncodingToString(this->current_encoding_); + throw ParquetException(ss.str()); + } + auto decoder = dynamic_cast*>(this->current_decoder_.get()); + const T* dictionary = nullptr; + decoder->GetDictionary(&dictionary, dictionary_length); + return reinterpret_cast(dictionary); +} + +template +int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records) { + if (num_records <= 0) { + return 0; + } + + Reserve(num_records); + + int64_t records_read = 0; + + do { + // Is there more data to read in this row group? + if (!this->HasNextInternal()) { + break; + } + + const int32_t batch_size = + std::min(clamp_to(num_records - records_read), + this->available_values_current_page()); + const auto decoded = + this->current_decoder_->Decode(values_.write_start(), batch_size); + CheckNumberDecoded(decoded, batch_size); + + values_.mark_values_as_written(decoded); + this->ConsumeBufferedValues(decoded); + + records_read += decoded; + } while (records_read < num_records); + + return records_read; +} + +template +void RequiredTypedRecordReader::Reset() { + if (values_written() > 0) { + values_.ResetValues(); + PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); + null_count_ = 0; + } +} + +template +void RequiredTypedRecordReader::SetPageReader(std::unique_ptr reader) { + this->pager_ = std::move(reader); + ResetDecoders(); +} + +template +void RequiredTypedRecordReader::DebugPrintState() { + const T* vals = reinterpret_cast(this->values()); + std::cout << "values: "; + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << vals[i] << " "; + } + std::cout << std::endl; +} + +// TODO(wesm): Implement these to some satisfaction +template <> +void RequiredTypedRecordReader::DebugPrintState() {} + /// In FLBARecordReader, we read fixed length byte array values. /// /// Unlike other fixed length types, the `values_` buffer is not used to store @@ -2641,29 +2784,43 @@ std::shared_ptr MakeByteArrayRecordReader( } // namespace +namespace { +template +std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* descr, + LevelInfo leaf_info, + MemoryPool* pool, + bool read_dense_for_nullable) { + if (descr->max_definition_level() == 0 && descr->max_repetition_level() == 0) { + return std::make_shared>(descr, pool); + } + return std::make_shared>(descr, leaf_info, pool, + read_dense_for_nullable); +} +} // namespace + std::shared_ptr RecordReader::Make( const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dictionary, bool read_dense_for_nullable, const std::shared_ptr<::arrow::DataType>& arrow_type) { switch (descr->physical_type()) { case Type::BOOLEAN: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::INT32: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::INT64: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::INT96: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::FLOAT: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::DOUBLE: - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); case Type::BYTE_ARRAY: { return MakeByteArrayRecordReader(descr, leaf_info, pool, read_dictionary, read_dense_for_nullable, arrow_type); From cb051b53e98aeaa4e665f56a603f7afcd73f829a Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 16 Jul 2026 14:28:31 +0200 Subject: [PATCH 12/33] Remove RecordReader protected members --- cpp/src/parquet/column_reader.cc | 110 ++++++++++++++++++++++++------- cpp/src/parquet/column_reader.h | 64 +++--------------- 2 files changed, 96 insertions(+), 78 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 436660fc59f..2242ea4aa8a 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1636,21 +1636,36 @@ class TypedRecordReader : public ColumnReaderImplBase, bool read_dense_for_nullable) : Base(descr, pool, LevelDecoder(descr->max_definition_level()), LevelDecoder(descr->max_repetition_level())), - values_(pool) { - leaf_info_ = leaf_info; - nullable_values_ = leaf_info_.HasNullableValues(); - at_record_start_ = true; - null_count_ = 0; - levels_written_ = 0; - levels_position_ = 0; - levels_capacity_ = 0; - read_dense_for_nullable_ = read_dense_for_nullable; - valid_bits_ = AllocateBuffer(pool); - def_levels_ = AllocateBuffer(pool); - rep_levels_ = AllocateBuffer(pool); + valid_bits_(AllocateBuffer(pool)), + values_(pool), + leaf_info_(leaf_info), + def_levels_(AllocateBuffer(pool)), + rep_levels_(AllocateBuffer(pool)), + nullable_values_(leaf_info.HasNullableValues()), + read_dense_for_nullable_(read_dense_for_nullable) { TypedRecordReader::Reset(); } + int16_t* def_levels() const final { + return reinterpret_cast(def_levels_->mutable_data()); + } + + int16_t* rep_levels() const final { + return reinterpret_cast(rep_levels_->mutable_data()); + } + + int64_t levels_position() const final { return levels_position_; } + + int64_t levels_written() const final { return levels_written_; } + + int64_t null_count() const final { return null_count_; } + + bool nullable_values() const final { return nullable_values_; } + + bool read_dictionary() const final { return read_dictionary_; } + + bool read_dense_for_nullable() const final { return read_dense_for_nullable_; } + uint8_t* values() const final { return reinterpret_cast(values_.data()); } int64_t values_written() const final { return values_.values_count(); } @@ -1758,9 +1773,54 @@ class TypedRecordReader : public ColumnReaderImplBase, void ResetValues(); + protected: + /// \brief Each bit corresponds to one element in 'values_' and specifies if it + /// is null or not null. + /// + /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. + std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; + bool read_dictionary_ = false; + private: ValuesBuffer values_; LevelInfo leaf_info_; + + /// \brief Buffer for definition levels. May contain more levels than + /// is actually read. This is because we read levels ahead to + /// figure out record boundaries for repeated fields. + /// For flat required fields, 'def_levels_' and 'rep_levels_' are not + /// populated. For non-repeated fields 'rep_levels_' is not populated. + /// 'def_levels_' and 'rep_levels_' must be of the same size if present. + std::shared_ptr<::arrow::ResizableBuffer> def_levels_; + /// \brief Buffer for repetition levels. Only populated for repeated + /// fields. + std::shared_ptr<::arrow::ResizableBuffer> rep_levels_; + + int64_t records_read_; + + int64_t null_count_ = 0; + + /// \brief Number of definition / repetition levels that have been written + /// internally in the reader. This may be larger than values_written() since + /// for repeated fields we need to look at the levels in advance to figure out + /// the record boundaries. + int64_t levels_written_ = 0; + /// \brief Position of the next level that should be consumed. + int64_t levels_position_ = 0; + int64_t levels_capacity_ = 0; + + /// \brief Indicates if we can have nullable values. Note that repeated fields + /// may or may not be nullable. + bool nullable_values_; + + bool at_record_start_ = true; + + // If true, we will not leave any space for the null values in the values_ + // vector or fill nulls values in BinaryRecordReader/DictionaryRecordReader. + // + // If read_dense_for_nullable_ is true, the BinaryRecordReader/DictionaryRecordReader + // might still populate the validity bitmap buffer. + bool read_dense_for_nullable_ = false; }; /************************************** @@ -2405,14 +2465,6 @@ class RequiredTypedRecordReader final : public ColumnReaderImplBasemax_definition_level()), LevelDecoder(descr->max_repetition_level())), values_(pool) { - nullable_values_ = false; - at_record_start_ = true; - null_count_ = 0; - levels_written_ = 0; - levels_position_ = 0; - levels_capacity_ = 0; - read_dense_for_nullable_ = false; - valid_bits_ = AllocateBuffer(pool); RequiredTypedRecordReader::Reset(); } @@ -2420,6 +2472,22 @@ class RequiredTypedRecordReader final : public ColumnReaderImplBase void RequiredTypedRecordReader::Reset() { if (values_written() > 0) { values_.ResetValues(); - PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); - null_count_ = 0; } } diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index 31a75d5b438..c23d15500d3 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -366,14 +366,10 @@ class PARQUET_EXPORT RecordReader { virtual const void* ReadDictionary(int32_t* dictionary_length) = 0; /// \brief Decoded definition levels - int16_t* def_levels() const { - return reinterpret_cast(def_levels_->mutable_data()); - } + virtual int16_t* def_levels() const = 0; /// \brief Decoded repetition levels - int16_t* rep_levels() const { - return reinterpret_cast(rep_levels_->mutable_data()); - } + virtual int16_t* rep_levels() const = 0; /// \brief Decoded values, including nulls, if any /// FLBA and ByteArray types do not use this array and read into their own @@ -389,71 +385,27 @@ class PARQUET_EXPORT RecordReader { /// \brief Number of definition / repetition levels (from those that have /// been decoded) that have been consumed inside the reader. - int64_t levels_position() const { return levels_position_; } + virtual int64_t levels_position() const = 0; /// \brief Number of definition / repetition levels that have been written /// internally in the reader. This may be larger than values_written() because /// for repeated fields we need to look at the levels in advance to figure out /// the record boundaries. - int64_t levels_written() const { return levels_written_; } + virtual int64_t levels_written() const = 0; /// \brief Number of nulls in the leaf that we have read so far into the /// values vector. This is only valid when !read_dense_for_nullable(). When /// read_dense_for_nullable() it will always be 0. - int64_t null_count() const { return null_count_; } + virtual int64_t null_count() const = 0; /// \brief True if the leaf values are nullable - bool nullable_values() const { return nullable_values_; } + virtual bool nullable_values() const = 0; /// \brief True if reading directly as Arrow dictionary-encoded - bool read_dictionary() const { return read_dictionary_; } + virtual bool read_dictionary() const = 0; /// \brief True if reading dense for nullable columns. - bool read_dense_for_nullable() const { return read_dense_for_nullable_; } - - protected: - /// \brief Indicates if we can have nullable values. Note that repeated fields - /// may or may not be nullable. - bool nullable_values_; - - bool at_record_start_; - int64_t records_read_; - - int64_t null_count_; - - /// \brief Each bit corresponds to one element in 'values_' and specifies if it - /// is null or not null. - /// - /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. - std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; - - /// \brief Buffer for definition levels. May contain more levels than - /// is actually read. This is because we read levels ahead to - /// figure out record boundaries for repeated fields. - /// For flat required fields, 'def_levels_' and 'rep_levels_' are not - /// populated. For non-repeated fields 'rep_levels_' is not populated. - /// 'def_levels_' and 'rep_levels_' must be of the same size if present. - std::shared_ptr<::arrow::ResizableBuffer> def_levels_; - /// \brief Buffer for repetition levels. Only populated for repeated - /// fields. - std::shared_ptr<::arrow::ResizableBuffer> rep_levels_; - - /// \brief Number of definition / repetition levels that have been written - /// internally in the reader. This may be larger than values_written() since - /// for repeated fields we need to look at the levels in advance to figure out - /// the record boundaries. - int64_t levels_written_; - /// \brief Position of the next level that should be consumed. - int64_t levels_position_; - int64_t levels_capacity_; - - bool read_dictionary_ = false; - // If true, we will not leave any space for the null values in the values_ - // vector or fill nulls values in BinaryRecordReader/DictionaryRecordReader. - // - // If read_dense_for_nullable_ is true, the BinaryRecordReader/DictionaryRecordReader - // might still populate the validity bitmap buffer. - bool read_dense_for_nullable_ = false; + virtual bool read_dense_for_nullable() const = 0; }; class BinaryRecordReader : virtual public RecordReader { From a903935e488fb66b90659b1d398544dee7ea8f78 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 16 Jul 2026 17:40:34 +0200 Subject: [PATCH 13/33] Fully use RequiredRecordReader --- cpp/src/parquet/column_reader.cc | 617 ++++++++++++++++++++----------- 1 file changed, 400 insertions(+), 217 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 2242ea4aa8a..491957ece48 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1170,6 +1170,10 @@ class ColumnReaderImplBase { int64_t Skip(int64_t num_values_to_skip); }; +/***************************************** + * ColumnReaderImplBase Implementation * + *****************************************/ + template int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* out) { int64_t num_decoded = current_decoder_->Decode(out, static_cast(batch_size)); @@ -1621,14 +1625,44 @@ namespace internal { namespace { +/// Whether values of type T can be printed to `std::cout`. +template +concept can_cout = requires(std::ostream& os, const T& value) { + os << value; +}; // NOLINT(readability/braces) + +/********************* + * ReadValuesHooks * + *********************/ + +/// Hooks a record reader calls to materialize decoded values. +/// +/// The default implementations in `TypedRecordReader` and +/// `RequiredTypedRecordReader` decode into the `values_` buffer. Leaf readers +/// for BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY columns override them to decode +/// directly into Arrow builders instead. +class ReadValuesHooks { + public: + virtual ~ReadValuesHooks() = default; + + virtual void ReserveValuesAndIsValid(int64_t extra_values) = 0; + + /// Decode `values_to_read` non-null values + virtual void ReadValuesDense(int64_t values_to_read) = 0; + + /// Decode a batch of `values_with_nulls` value slots, `null_count` + /// of which are null + virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) = 0; +}; + /*********************** * TypedRecordReader * ***********************/ -template > +template class TypedRecordReader : public ColumnReaderImplBase, - virtual public RecordReader { + virtual public RecordReader, + public ReadValuesHooks { public: using T = typename DType::c_type; using Base = ColumnReaderImplBase; @@ -1662,7 +1696,7 @@ class TypedRecordReader : public ColumnReaderImplBase, bool nullable_values() const final { return nullable_values_; } - bool read_dictionary() const final { return read_dictionary_; } + bool read_dictionary() const final { return kReadDictionary; } bool read_dense_for_nullable() const final { return read_dense_for_nullable_; } @@ -1726,7 +1760,7 @@ class TypedRecordReader : public ColumnReaderImplBase, void ReserveLevels(int64_t extra_levels); - virtual void ReserveValuesAndIsValid(int64_t extra_values); + void ReserveValuesAndIsValid(int64_t extra_values) override; void Reset() override; @@ -1739,9 +1773,9 @@ class TypedRecordReader : public ColumnReaderImplBase, // Dictionary decoders must be reset when advancing row groups void ResetDecoders() { this->decoders_.clear(); } - virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count); + void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) override; - virtual void ReadValuesDense(int64_t values_to_read); + void ReadValuesDense(int64_t values_to_read) override; // Reads repeated records and returns number of records read. Fills in // values_to_read and null_count. @@ -1753,10 +1787,6 @@ class TypedRecordReader : public ColumnReaderImplBase, int64_t ReadOptionalRecords(int64_t num_records, int64_t* values_to_read, int64_t* null_count); - // Reads required records and returns number of records read. Fills in - // values_to_read. - int64_t ReadRequiredRecords(int64_t num_records, int64_t* values_to_read); - // Reads dense for optional records. First it figures out how many values to // read. void ReadDenseForOptional(int64_t start_levels_position, int64_t* values_to_read); @@ -1779,7 +1809,6 @@ class TypedRecordReader : public ColumnReaderImplBase, /// /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; - bool read_dictionary_ = false; private: ValuesBuffer values_; @@ -1827,8 +1856,8 @@ class TypedRecordReader : public ColumnReaderImplBase, * TypedRecordReader Implementation * **************************************/ -template -const void* TypedRecordReader::ReadDictionary( +template +const void* TypedRecordReader::ReadDictionary( int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; @@ -1849,8 +1878,9 @@ const void* TypedRecordReader::ReadDictionary( return reinterpret_cast(dictionary); } -template -int64_t TypedRecordReader::ReadRecords(int64_t num_records) { +template +int64_t TypedRecordReader::ReadRecords( + int64_t num_records) { if (num_records == 0) return 0; // Delimit records, then read values at the end int64_t records_read = 0; @@ -1918,8 +1948,8 @@ int64_t TypedRecordReader::ReadRecords(int64_t num_records) return records_read; } -template -void TypedRecordReader::ThrowAwayLevels( +template +void TypedRecordReader::ThrowAwayLevels( int64_t start_levels_position) { ARROW_DCHECK_LE(levels_position_, levels_written_); ARROW_DCHECK_LE(start_levels_position, levels_position_); @@ -1951,8 +1981,9 @@ void TypedRecordReader::ThrowAwayLevels( levels_capacity_ -= gap; } -template -int64_t TypedRecordReader::SkipRecordsInBufferNonRepeated( +template +int64_t +TypedRecordReader::SkipRecordsInBufferNonRepeated( int64_t num_records) { ARROW_DCHECK_EQ(this->max_rep_level(), 0); if (!this->has_values_to_process() || num_records == 0) return 0; @@ -1985,8 +2016,9 @@ int64_t TypedRecordReader::SkipRecordsInBufferNonRepeated( return skipped_records; } -template -int64_t TypedRecordReader::DelimitAndSkipRecordsInBuffer( +template +int64_t +TypedRecordReader::DelimitAndSkipRecordsInBuffer( int64_t num_records) { if (num_records == 0) return 0; // Look at the buffered levels, delimit them based on @@ -2006,8 +2038,9 @@ int64_t TypedRecordReader::DelimitAndSkipRecordsInBuffer( return skipped_records; } -template -int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_records) { +template +int64_t TypedRecordReader::SkipRecordsRepeated( + int64_t num_records) { ARROW_DCHECK_GT(this->max_rep_level(), 0); int64_t skipped_records = 0; @@ -2073,8 +2106,9 @@ int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_ return skipped_records; } -template -void TypedRecordReader::ReadAndThrowAwayValues(int64_t num_values) { +template +void TypedRecordReader::ReadAndThrowAwayValues( + int64_t num_values) { const int64_t values_read = this->current_decoder_.Skip(num_values); if (values_read < num_values) { std::stringstream ss; @@ -2083,8 +2117,9 @@ void TypedRecordReader::ReadAndThrowAwayValues(int64_t num_ } } -template -int64_t TypedRecordReader::SkipRecords(int64_t num_records) { +template +int64_t TypedRecordReader::SkipRecords( + int64_t num_records) { if (num_records == 0) return 0; // Top level required field. Number of records equals to number of levels, @@ -2111,9 +2146,9 @@ int64_t TypedRecordReader::SkipRecords(int64_t num_records) return skipped_records; } -template +template std::shared_ptr -TypedRecordReader::ReleaseIsValid() { +TypedRecordReader::ReleaseIsValid() { if (nullable_values()) { const auto bit_count = bit_util::BytesForBits(values_written()); auto result = valid_bits_; @@ -2126,9 +2161,9 @@ TypedRecordReader::ReleaseIsValid() { } } -template -int64_t TypedRecordReader::DelimitRecords(int64_t num_records, - int64_t* values_seen) { +template +int64_t TypedRecordReader::DelimitRecords( + int64_t num_records, int64_t* values_seen) { if (ARROW_PREDICT_FALSE(num_records == 0 || levels_position_ == levels_written_)) { *values_seen = 0; return 0; @@ -2186,14 +2221,15 @@ int64_t TypedRecordReader::DelimitRecords(int64_t num_recor return records_read; } -template -void TypedRecordReader::Reserve(int64_t capacity) { +template +void TypedRecordReader::Reserve(int64_t capacity) { ReserveLevels(capacity); ReserveValuesAndIsValid(capacity); } -template -void TypedRecordReader::ReserveLevels(int64_t extra_levels) { +template +void TypedRecordReader::ReserveLevels( + int64_t extra_levels) { if (this->max_def_level() > 0) { const int64_t new_levels_capacity = compute_capacity_pow2(levels_capacity_, levels_written_, extra_levels); @@ -2214,8 +2250,8 @@ void TypedRecordReader::ReserveLevels(int64_t extra_levels) } } -template -void TypedRecordReader::ReserveValuesAndIsValid( +template +void TypedRecordReader::ReserveValuesAndIsValid( int64_t extra_values) { values_.ReserveValues(extra_values); if (nullable_values() && !read_dense_for_nullable_) { @@ -2231,8 +2267,8 @@ void TypedRecordReader::ReserveValuesAndIsValid( } } -template -void TypedRecordReader::Reset() { +template +void TypedRecordReader::Reset() { ResetValues(); if (levels_written_ > 0) { @@ -2243,17 +2279,17 @@ void TypedRecordReader::Reset() { // Call Finish on the binary builders to reset them } -template -void TypedRecordReader::SetPageReader( +template +void TypedRecordReader::SetPageReader( std::unique_ptr reader) { at_record_start_ = true; this->pager_ = std::move(reader); ResetDecoders(); } -template -void TypedRecordReader::ReadValuesSpaced(int64_t values_with_nulls, - int64_t null_count) { +template +void TypedRecordReader::ReadValuesSpaced( + int64_t values_with_nulls, int64_t null_count) { uint8_t* valid_bits = valid_bits_->mutable_data(); const int64_t valid_bits_offset = values_written(); @@ -2263,15 +2299,16 @@ void TypedRecordReader::ReadValuesSpaced(int64_t values_wit CheckNumberDecoded(num_decoded, values_with_nulls); } -template -void TypedRecordReader::ReadValuesDense(int64_t values_to_read) { +template +void TypedRecordReader::ReadValuesDense( + int64_t values_to_read) { int64_t num_decoded = this->current_decoder_->Decode(values_.write_start(), static_cast(values_to_read)); CheckNumberDecoded(num_decoded, values_to_read); } -template -int64_t TypedRecordReader::ReadRepeatedRecords( +template +int64_t TypedRecordReader::ReadRepeatedRecords( int64_t num_records, int64_t* values_to_read, int64_t* null_count) { const int64_t start_levels_position = levels_position_; // Note that repeated records may be required or nullable. If they have @@ -2292,8 +2329,8 @@ int64_t TypedRecordReader::ReadRepeatedRecords( return records_read; } -template -int64_t TypedRecordReader::ReadOptionalRecords( +template +int64_t TypedRecordReader::ReadOptionalRecords( int64_t num_records, int64_t* values_to_read, int64_t* null_count) { const int64_t start_levels_position = levels_position_; // No repetition levels, skip delimiting logic. Each level represents a @@ -2315,16 +2352,8 @@ int64_t TypedRecordReader::ReadOptionalRecords( return records_read; } -template -int64_t TypedRecordReader::ReadRequiredRecords( - int64_t num_records, int64_t* values_to_read) { - *values_to_read = num_records; - ReadValuesDense(*values_to_read); - return num_records; -} - -template -void TypedRecordReader::ReadDenseForOptional( +template +void TypedRecordReader::ReadDenseForOptional( int64_t start_levels_position, int64_t* values_to_read) { // levels_position_ must already be incremented based on number of records // read. @@ -2337,9 +2366,10 @@ void TypedRecordReader::ReadDenseForOptional( ReadValuesDense(*values_to_read); } -template -void TypedRecordReader::ReadSpacedForOptionalOrRepeated( - int64_t start_levels_position, int64_t* values_to_read, int64_t* null_count) { +template +void TypedRecordReader:: + ReadSpacedForOptionalOrRepeated(int64_t start_levels_position, + int64_t* values_to_read, int64_t* null_count) { // levels_position_ must already be incremented based on number of records // read. ARROW_DCHECK_GE(levels_position_, start_levels_position); @@ -2357,8 +2387,9 @@ void TypedRecordReader::ReadSpacedForOptionalOrRepeated( ReadValuesSpaced(validity_io.values_read, *null_count); } -template -int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { +template +int64_t TypedRecordReader::ReadRecordData( + int64_t num_records) { // Conservative upper bound const int64_t possible_num_values = std::max(num_records, levels_written_ - levels_position_); @@ -2382,7 +2413,9 @@ int64_t TypedRecordReader::ReadRecordData(int64_t num_recor records_read = ReadOptionalRecords(num_records, &values_to_read, &null_count); } else { ARROW_DCHECK(!nullable_values()); - records_read = ReadRequiredRecords(num_records, &values_to_read); + values_to_read = num_records; + ReadValuesDense(values_to_read); + records_read = num_records; // We don't need to update null_count, since it is 0. } @@ -2409,14 +2442,12 @@ int64_t TypedRecordReader::ReadRecordData(int64_t num_recor return records_read; } -template -void TypedRecordReader::DebugPrintState() { +template +void TypedRecordReader::DebugPrintState() { const int16_t* def_levels = this->def_levels(); const int16_t* rep_levels = this->rep_levels(); const int64_t total_levels_read = levels_position_; - const T* vals = reinterpret_cast(this->values()); - if (leaf_info_.def_level > 0) { std::cout << "def levels: "; for (int64_t i = 0; i < total_levels_read; ++i) { @@ -2434,14 +2465,21 @@ void TypedRecordReader::DebugPrintState() { } std::cout << "values: "; - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << vals[i] << " "; + if constexpr (can_cout) { + const T* vals = reinterpret_cast(this->values()); + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << vals[i] << " "; + } + } else { + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << "? "; + } } std::cout << std::endl; } -template -void TypedRecordReader::ResetValues() { +template +void TypedRecordReader::ResetValues() { if (values_written() > 0) { values_.ResetValues(); PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); @@ -2453,13 +2491,15 @@ void TypedRecordReader::ResetValues() { * RequiredTypedRecordReader * *******************************/ -template -class RequiredTypedRecordReader final : public ColumnReaderImplBase, - public RecordReader { +template , + bool kReadDictionary = false> +class RequiredTypedRecordReader : public ColumnReaderImplBase, + virtual public RecordReader, + public ReadValuesHooks { public: using T = typename DType::c_type; using Base = ColumnReaderImplBase; - using ValuesBuffer = ReadValuesBuffer; RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool) : Base(descr, pool, LevelDecoder(descr->max_definition_level()), @@ -2484,36 +2524,51 @@ class RequiredTypedRecordReader final : public ColumnReaderImplBaseSkip(num_records); } + int64_t SkipRecords(int64_t num_records) final { return this->Skip(num_records); } - std::shared_ptr ReleaseValues() override { + std::shared_ptr ReleaseValues() final { return values_.ReleaseValues(this->pool_); } - std::shared_ptr ReleaseIsValid() override { return nullptr; } + std::shared_ptr ReleaseIsValid() final { return nullptr; } - void Reserve(int64_t extra_values) override { values_.ReserveValues(extra_values); } + void Reserve(int64_t extra_values) final { ReserveValuesAndIsValid(extra_values); } - void Reset() override; + void Reset() final; - void SetPageReader(std::unique_ptr reader) override; + void SetPageReader(std::unique_ptr reader) final; - bool HasMoreData() const override { return this->pager_ != nullptr; } + bool HasMoreData() const final { return this->pager_ != nullptr; } - const ColumnDescriptor* descr() const override { return this->descr_; } + const ColumnDescriptor* descr() const final { return this->descr_; } // Dictionary decoders must be reset when advancing row groups void ResetDecoders() { this->decoders_.clear(); } - void DebugPrintState() override; + void DebugPrintState() final; + + protected: + // ReadValuesHooks default implementations, decoding into `values_`. + // Leaf classes for binary types override these to decode into their own + // Arrow builders. + void ReserveValuesAndIsValid(int64_t extra_values) override { + values_.ReserveValues(extra_values); + } + + void ReadValuesDense(int64_t values_to_read) override; + + // A required column never contains nulls, so values are never read spaced. + void ReadValuesSpaced(int64_t /*values_with_nulls*/, int64_t /*null_count*/) override { + throw ParquetException("Spaced values read in required column"); + } private: ValuesBuffer values_; @@ -2523,8 +2578,10 @@ class RequiredTypedRecordReader final : public ColumnReaderImplBase -const void* RequiredTypedRecordReader::ReadDictionary(int32_t* dictionary_length) { +template +const void* +RequiredTypedRecordReader::ReadDictionary( + int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; return nullptr; @@ -2544,8 +2601,9 @@ const void* RequiredTypedRecordReader::ReadDictionary(int32_t* dictionary return reinterpret_cast(dictionary); } -template -int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records) { +template +int64_t RequiredTypedRecordReader::ReadRecords( + int64_t num_records) { if (num_records <= 0) { return 0; } @@ -2563,65 +2621,119 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records) { const int32_t batch_size = std::min(clamp_to(num_records - records_read), this->available_values_current_page()); - const auto decoded = - this->current_decoder_->Decode(values_.write_start(), batch_size); - CheckNumberDecoded(decoded, batch_size); + ReadValuesDense(batch_size); - values_.mark_values_as_written(decoded); - this->ConsumeBufferedValues(decoded); + values_.mark_values_as_written(batch_size); + this->ConsumeBufferedValues(batch_size); - records_read += decoded; + records_read += batch_size; } while (records_read < num_records); return records_read; } -template -void RequiredTypedRecordReader::Reset() { +template +void RequiredTypedRecordReader::ReadValuesDense( + int64_t values_to_read) { + const int64_t num_decoded = this->current_decoder_->Decode( + values_.write_start(), static_cast(values_to_read)); + CheckNumberDecoded(num_decoded, values_to_read); +} + +template +void RequiredTypedRecordReader::Reset() { if (values_written() > 0) { values_.ResetValues(); } } -template -void RequiredTypedRecordReader::SetPageReader(std::unique_ptr reader) { +template +void RequiredTypedRecordReader::SetPageReader( + std::unique_ptr reader) { this->pager_ = std::move(reader); ResetDecoders(); } -template -void RequiredTypedRecordReader::DebugPrintState() { - const T* vals = reinterpret_cast(this->values()); +template +void RequiredTypedRecordReader::DebugPrintState() { std::cout << "values: "; - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << vals[i] << " "; + if constexpr (can_cout) { + const T* vals = reinterpret_cast(this->values()); + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << vals[i] << " "; + } + } else { + for (int64_t i = 0; i < this->values_written(); ++i) { + std::cout << "? "; + } } std::cout << std::endl; } -// TODO(wesm): Implement these to some satisfaction -template <> -void RequiredTypedRecordReader::DebugPrintState() {} +/************************ + * record_reader_base * + ************************/ + +template +struct record_reader_base; + +template +struct record_reader_base { + using c_type = typename DType::c_type; + using ValuesBuffer = ReadValuesNoBuffer; + using type = RequiredTypedRecordReader; +}; + +template +struct record_reader_base { + using c_type = typename DType::c_type; + using ValuesBuffer = ReadValuesNoBuffer; + using type = TypedRecordReader; +}; + +template +using record_reader_base_t = + typename record_reader_base::type; -/// In FLBARecordReader, we read fixed length byte array values. +/********************** + * FLBARecordReader * + **********************/ + +/// Reads fixed length byte array values into a FixedSizeBinaryBuilder. /// -/// Unlike other fixed length types, the `values_` buffer is not used to store -/// values, instead we use `data_builder_` to store the values, and `null_bitmap_builder_` -/// is used to store the null bitmap. +/// `kRequired` selects the base class: RequiredTypedRecordReader for +/// required (non-nullable, non-repeated) columns, TypedRecordReader +/// otherwise. /// -/// The `values_` buffer is used to store the temporary values for `Decode`, and it would -/// be Reset after each `Decode` call. The `valid_bits_` buffer is never used. -class FLBARecordReader final - : public TypedRecordReader>, - virtual public BinaryRecordReader { +/// Values are decoded directly into `array_builder_`; the `values_` buffer +/// of the base class is a ReadValuesNoBuffer and only tracks the number of +/// values written. The `valid_bits_` buffer, if any, is consumed by each +/// spaced decode. +template +class FLBARecordReader final : public record_reader_base_t, + virtual public BinaryRecordReader { public: + using Base = record_reader_base_t; + FLBARecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) - : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), - byte_width_(descr_->type_length()), + requires(!kRequired) + : Base(descr, leaf_info, pool, read_dense_for_nullable), + byte_width_(descr->type_length()), + type_(::arrow::fixed_size_binary(byte_width_)), + array_builder_(type_, pool) { + ARROW_DCHECK_EQ(descr->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); + } + + FLBARecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) + requires(kRequired) + : Base(descr, pool), + byte_width_(descr->type_length()), type_(::arrow::fixed_size_binary(byte_width_)), array_builder_(type_, pool) { - ARROW_DCHECK_EQ(descr_->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); + ARROW_DCHECK_EQ(descr->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); + ARROW_DCHECK_EQ(descr->max_definition_level(), 0); + ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); } ::arrow::ArrayVector GetBuilderChunks() override { @@ -2629,8 +2741,9 @@ class FLBARecordReader final return ::arrow::ArrayVector{std::move(chunk)}; } + protected: void ReserveValuesAndIsValid(int64_t extra_values) override { - TypedRecordReader::ReserveValuesAndIsValid(extra_values); + Base::ReserveValuesAndIsValid(extra_values); PARQUET_THROW_NOT_OK(array_builder_.Reserve(extra_values)); } @@ -2638,15 +2751,22 @@ class FLBARecordReader final int64_t num_decoded = this->current_decoder_->DecodeArrowNonNull( static_cast(values_to_read), &array_builder_); CheckNumberDecoded(num_decoded, values_to_read); - ResetValues(); + if constexpr (!kRequired) { + this->ResetValues(); + } } void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - int64_t num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written(), &array_builder_); - CheckNumberDecoded(num_decoded, values_to_read - null_count); - ResetValues(); + if constexpr (kRequired) { + // A required column never contains nulls: the base implementation throws. + Base::ReadValuesSpaced(values_to_read, null_count); + } else { + int64_t num_decoded = this->current_decoder_->DecodeArrow( + static_cast(values_to_read), static_cast(null_count), + this->valid_bits_->mutable_data(), this->values_written(), &array_builder_); + CheckNumberDecoded(num_decoded, values_to_read - null_count); + this->ResetValues(); + } } private: @@ -2655,47 +2775,69 @@ class FLBARecordReader final ::arrow::FixedSizeBinaryBuilder array_builder_; }; -/// ByteArrayRecordReader reads variable length byte array values. +/*************************** + * ByteArrayRecordReader * + ***************************/ + +/// Create the Arrow builder for reading a Parquet BYTE_ARRAY column as the +/// given Arrow type (BINARY by default). +std::unique_ptr<::arrow::ArrayBuilder> MakeByteArrayBuilder( + const std::shared_ptr<::arrow::DataType>& arrow_type, ::arrow::MemoryPool* pool) { + auto arrow_binary_type = arrow_type ? arrow_type->id() : ::arrow::Type::BINARY; + switch (arrow_binary_type) { + case ::arrow::Type::BINARY: + return std::make_unique<::arrow::BinaryBuilder>(pool); + case ::arrow::Type::STRING: + return std::make_unique<::arrow::StringBuilder>(pool); + case ::arrow::Type::LARGE_BINARY: + return std::make_unique<::arrow::LargeBinaryBuilder>(pool); + case ::arrow::Type::LARGE_STRING: + return std::make_unique<::arrow::LargeStringBuilder>(pool); + case ::arrow::Type::BINARY_VIEW: + return std::make_unique<::arrow::BinaryViewBuilder>(pool); + case ::arrow::Type::STRING_VIEW: + return std::make_unique<::arrow::StringViewBuilder>(pool); + default: + throw ParquetException("cannot read Parquet BYTE_ARRAY as Arrow " + + arrow_type->ToString()); + } +} + +/// Reads variable length byte array values into a chunked binary builder. +/// +/// `kRequired` selects the base class: RequiredTypedRecordReader for +/// required (non-nullable, non-repeated) columns, TypedRecordReader +/// otherwise. /// /// It only calls `DecodeArrowNonNull` and `DecodeArrow` to read values, and /// `Decode` and `DecodeSpaced` are not used. /// /// The `values_` buffers are never used, and the `accumulator_` /// is used to store the values. +template class ByteArrayChunkedRecordReader final - : public TypedRecordReader>, + : public record_reader_base_t, virtual public BinaryRecordReader { public: + using Base = record_reader_base_t; + ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable, const std::shared_ptr<::arrow::DataType>& arrow_type) - : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable) { - ARROW_DCHECK_EQ(descr_->physical_type(), Type::BYTE_ARRAY); - auto arrow_binary_type = arrow_type ? arrow_type->id() : ::arrow::Type::BINARY; - switch (arrow_binary_type) { - case ::arrow::Type::BINARY: - accumulator_.builder = std::make_unique<::arrow::BinaryBuilder>(pool); - break; - case ::arrow::Type::STRING: - accumulator_.builder = std::make_unique<::arrow::StringBuilder>(pool); - break; - case ::arrow::Type::LARGE_BINARY: - accumulator_.builder = std::make_unique<::arrow::LargeBinaryBuilder>(pool); - break; - case ::arrow::Type::LARGE_STRING: - accumulator_.builder = std::make_unique<::arrow::LargeStringBuilder>(pool); - break; - case ::arrow::Type::BINARY_VIEW: - accumulator_.builder = std::make_unique<::arrow::BinaryViewBuilder>(pool); - break; - case ::arrow::Type::STRING_VIEW: - accumulator_.builder = std::make_unique<::arrow::StringViewBuilder>(pool); - break; - default: - throw ParquetException("cannot read Parquet BYTE_ARRAY as Arrow " + - arrow_type->ToString()); - } + requires(!kRequired) + : Base(descr, leaf_info, pool, read_dense_for_nullable) { + ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); + accumulator_.builder = MakeByteArrayBuilder(arrow_type, pool); + } + + ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, + const std::shared_ptr<::arrow::DataType>& arrow_type) + requires(kRequired) + : Base(descr, pool) { + ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); + ARROW_DCHECK_EQ(descr->max_definition_level(), 0); + ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); + accumulator_.builder = MakeByteArrayBuilder(arrow_type, pool); } ::arrow::ArrayVector GetBuilderChunks() override { @@ -2709,8 +2851,9 @@ class ByteArrayChunkedRecordReader final return result; } + protected: void ReserveValuesAndIsValid(int64_t extra_values) override { - TypedRecordReader::ReserveValuesAndIsValid(extra_values); + Base::ReserveValuesAndIsValid(extra_values); PARQUET_THROW_NOT_OK(accumulator_.builder->Reserve(extra_values)); } @@ -2718,15 +2861,22 @@ class ByteArrayChunkedRecordReader final int64_t num_decoded = this->current_decoder_->DecodeArrowNonNull( static_cast(values_to_read), &accumulator_); CheckNumberDecoded(num_decoded, values_to_read); - ResetValues(); + if constexpr (!kRequired) { + this->ResetValues(); + } } void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - int64_t num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written(), &accumulator_); - CheckNumberDecoded(num_decoded, values_to_read - null_count); - ResetValues(); + if constexpr (kRequired) { + // A required column never contains nulls: the base implementation throws. + Base::ReadValuesSpaced(values_to_read, null_count); + } else { + int64_t num_decoded = this->current_decoder_->DecodeArrow( + static_cast(values_to_read), static_cast(null_count), + this->valid_bits_->mutable_data(), this->values_written(), &accumulator_); + CheckNumberDecoded(num_decoded, values_to_read - null_count); + this->ResetValues(); + } } private: @@ -2734,21 +2884,36 @@ class ByteArrayChunkedRecordReader final typename EncodingTraits::Accumulator accumulator_; }; -/// ByteArrayDictionaryRecordReader reads into ::arrow::dictionary(index: int32, -/// values: binary). +/// Reads byte array values into ::arrow::dictionary(index: int32, values: binary). /// -/// If underlying column is dictionary encoded, it will call `DecodeIndices` to read, -/// otherwise it will call `DecodeArrowNonNull` to read. +/// `kRequired` selects the base class: RequiredTypedRecordReader for +/// required (non-nullable, non-repeated) columns, TypedRecordReader +/// otherwise. +/// +/// If the underlying column is dictionary encoded, it will call +/// `DecodeIndices` to read, otherwise it will call `DecodeArrowNonNull` to +/// read. +template class ByteArrayDictionaryRecordReader final - : public TypedRecordReader>, + : public record_reader_base_t, virtual public DictionaryRecordReader { public: + using Base = record_reader_base_t; + ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) - : TypedRecordReader(descr, leaf_info, pool, read_dense_for_nullable), - builder_(pool) { - this->read_dictionary_ = true; + requires(!kRequired) + : Base(descr, leaf_info, pool, read_dense_for_nullable), builder_(pool) { + ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); + } + + ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, + ::arrow::MemoryPool* pool) + requires(kRequired) + : Base(descr, pool), builder_(pool) { + ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); + ARROW_DCHECK_EQ(descr->max_definition_level(), 0); + ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); } std::shared_ptr<::arrow::ChunkedArray> GetResult() override { @@ -2758,6 +2923,7 @@ class ByteArrayDictionaryRecordReader final return std::make_shared<::arrow::ChunkedArray>(std::move(result), builder_.type()); } + protected: void FlushBuilder() { if (builder_.length() > 0) { std::shared_ptr<::arrow::Array> chunk; @@ -2783,7 +2949,7 @@ class ByteArrayDictionaryRecordReader final void ReadValuesDense(int64_t values_to_read) override { int64_t num_decoded = 0; - if (current_encoding_ == Encoding::RLE_DICTIONARY) { + if (this->current_encoding_ == Encoding::RLE_DICTIONARY) { MaybeWriteNewDictionary(); auto decoder = dynamic_cast(this->current_decoder_.get()); num_decoded = decoder->DecodeIndices(static_cast(values_to_read), &builder_); @@ -2791,27 +2957,34 @@ class ByteArrayDictionaryRecordReader final num_decoded = this->current_decoder_->DecodeArrowNonNull( static_cast(values_to_read), &builder_); } - // Flush values since they have been copied into the builder - ResetValues(); + if constexpr (!kRequired) { + // Flush values since they have been copied into the builder + this->ResetValues(); + } CheckNumberDecoded(num_decoded, values_to_read); } void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - int64_t num_decoded = 0; - if (current_encoding_ == Encoding::RLE_DICTIONARY) { - MaybeWriteNewDictionary(); - 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_); + if constexpr (kRequired) { + // A required column never contains nulls: the base implementation throws. + Base::ReadValuesSpaced(values_to_read, null_count); } else { - num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - valid_bits_->mutable_data(), values_written(), &builder_); + int64_t num_decoded = 0; + if (this->current_encoding_ == Encoding::RLE_DICTIONARY) { + MaybeWriteNewDictionary(); + auto decoder = dynamic_cast(this->current_decoder_.get()); + num_decoded = decoder->DecodeIndicesSpaced( + static_cast(values_to_read), static_cast(null_count), + this->valid_bits_->mutable_data(), this->values_written(), &builder_); + } else { + num_decoded = this->current_decoder_->DecodeArrow( + static_cast(values_to_read), static_cast(null_count), + this->valid_bits_->mutable_data(), this->values_written(), &builder_); + } + ARROW_DCHECK_EQ(num_decoded, values_to_read - null_count); + // Flush values since they have been copied into the builder + this->ResetValues(); } - ARROW_DCHECK_EQ(num_decoded, values_to_read - null_count); - // Flush values since they have been copied into the builder - ResetValues(); } private: @@ -2821,30 +2994,29 @@ class ByteArrayDictionaryRecordReader final std::vector> result_chunks_; }; -// TODO(wesm): Implement these to some satisfaction -template <> -void TypedRecordReader::DebugPrintState() {} - -template <> -void TypedRecordReader< - ByteArrayType, - ReadValuesNoBuffer>::DebugPrintState() {} - -template <> -void TypedRecordReader>::DebugPrintState() { -} - std::shared_ptr MakeByteArrayRecordReader( const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dictionary, bool read_dense_for_nullable, const std::shared_ptr<::arrow::DataType>& arrow_type) { + const bool required = + descr->max_definition_level() == 0 && descr->max_repetition_level() == 0; if (read_dictionary) { - return std::make_shared(descr, leaf_info, pool, - read_dense_for_nullable); + if (required) { + using RequiredReader = ByteArrayDictionaryRecordReader; + + return std::make_shared(descr, pool); + } + using Reader = ByteArrayDictionaryRecordReader; + return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } else { - return std::make_shared( - descr, leaf_info, pool, read_dense_for_nullable, arrow_type); + if (required) { + using RequiredReader = ByteArrayChunkedRecordReader; + + return std::make_shared(descr, pool, arrow_type); + } + using Reader = ByteArrayChunkedRecordReader; + return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable, + arrow_type); } } @@ -2857,10 +3029,21 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* MemoryPool* pool, bool read_dense_for_nullable) { if (descr->max_definition_level() == 0 && descr->max_repetition_level() == 0) { - return std::make_shared>(descr, pool); + if constexpr (std::is_same_v) { + using FLBAReader = FLBARecordReader; + return std::make_shared(descr, pool); + } else { + return std::make_shared>(descr, pool); + } + } + if constexpr (std::is_same_v) { + using FLBAReader = FLBARecordReader; + return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); + } else { + using c_type = typename DType::c_type; + using Reader = TypedRecordReader, false>; + return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } - return std::make_shared>(descr, leaf_info, pool, - read_dense_for_nullable); } } // namespace @@ -2892,8 +3075,8 @@ std::shared_ptr RecordReader::Make( read_dense_for_nullable, arrow_type); } case Type::FIXED_LEN_BYTE_ARRAY: - return std::make_shared(descr, leaf_info, pool, - read_dense_for_nullable); + return DispatchTypedRecordReader(descr, leaf_info, pool, + read_dense_for_nullable); default: { // PARQUET-1481: This can occur if the file is corrupt std::stringstream ss; From 09ce174b4cb80b6590ce73af4bbab743e778c1df Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 17 Jul 2026 11:25:00 +0200 Subject: [PATCH 14/33] Rename ReadValues -> ValueSink --- cpp/src/parquet/column_reader.cc | 235 ++++++++++++++++--------------- 1 file changed, 119 insertions(+), 116 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 491957ece48..3bfb2d2b36b 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -726,9 +726,9 @@ class SkippableTypedDecoder { } }; -/****************** - * BufferCursor * - ******************/ +/********************* + * ValueSinkCursor * + *********************/ inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { if (extra_size < 0) { @@ -747,7 +747,7 @@ inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t ext return bit_util::NextPower2(target_size); } -class BufferCursor { +class ValueSinkCursor { public: int64_t capacity() const { return capacity_; } @@ -768,19 +768,19 @@ class BufferCursor { int64_t capacity_ = 0; }; -/********************** - * ReadValuesBuffer * - **********************/ +/********************* + * ValueSinkBuffer * + *********************/ template -class ReadValuesBuffer : private BufferCursor { +class ValueSinkBuffer : private ValueSinkCursor { public: using value_type = T; - using BufferCursor::capacity; - using BufferCursor::values_count; + using ValueSinkCursor::capacity; + using ValueSinkCursor::values_count; - explicit ReadValuesBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} + explicit ValueSinkBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} value_type* data() const { return values_->mutable_data_as(); } @@ -816,7 +816,7 @@ class ReadValuesBuffer : private BufferCursor { void ResetValues() { if (values_count() > 0) { PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); - BufferCursor::operator=({}); + ValueSinkCursor::operator=({}); } } @@ -833,41 +833,11 @@ class ReadValuesBuffer : private BufferCursor { } }; -/************************ - * ReadValuesNoBuffer * - ************************/ - -template -class ReadValuesNoBuffer : private BufferCursor { - public: - using value_type = T; - - using BufferCursor::capacity; - using BufferCursor::values_count; - - explicit ReadValuesNoBuffer(MemoryPool* /* pool */) {} - - value_type* data() const { return nullptr; } - - value_type* write_start() { return nullptr; } - - void mark_values_as_written(int64_t extra_values) { - set_values_count(values_count() + extra_values); - } - - std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { - return nullptr; - } - - void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } - - void ResetValues() { BufferCursor::operator=({}); } -}; - /************************* * StackedLevelDecoder * *************************/ +// TODO do we still need this? class StackedLevelDecoder : private LevelDecoder { public: using Base = LevelDecoder; @@ -895,7 +865,7 @@ class StackedLevelDecoder : private LevelDecoder { int64_t size() const { return position_; } private: - ReadValuesBuffer buffer_; + ValueSinkBuffer buffer_; int64_t position_ = 0; int64_t total_decoded() const { return buffer_.values_count(); } @@ -1659,7 +1629,7 @@ class ReadValuesHooks { * TypedRecordReader * ***********************/ -template +template class TypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader, public ReadValuesHooks { @@ -1811,7 +1781,7 @@ class TypedRecordReader : public ColumnReaderImplBase, std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; private: - ValuesBuffer values_; + ValueSink values_; LevelInfo leaf_info_; /// \brief Buffer for definition levels. May contain more levels than @@ -1856,8 +1826,8 @@ class TypedRecordReader : public ColumnReaderImplBase, * TypedRecordReader Implementation * **************************************/ -template -const void* TypedRecordReader::ReadDictionary( +template +const void* TypedRecordReader::ReadDictionary( int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; @@ -1878,8 +1848,8 @@ const void* TypedRecordReader::ReadDiction return reinterpret_cast(dictionary); } -template -int64_t TypedRecordReader::ReadRecords( +template +int64_t TypedRecordReader::ReadRecords( int64_t num_records) { if (num_records == 0) return 0; // Delimit records, then read values at the end @@ -1948,8 +1918,8 @@ int64_t TypedRecordReader::ReadRecords( return records_read; } -template -void TypedRecordReader::ThrowAwayLevels( +template +void TypedRecordReader::ThrowAwayLevels( int64_t start_levels_position) { ARROW_DCHECK_LE(levels_position_, levels_written_); ARROW_DCHECK_LE(start_levels_position, levels_position_); @@ -1981,9 +1951,9 @@ void TypedRecordReader::ThrowAwayLevels( levels_capacity_ -= gap; } -template +template int64_t -TypedRecordReader::SkipRecordsInBufferNonRepeated( +TypedRecordReader::SkipRecordsInBufferNonRepeated( int64_t num_records) { ARROW_DCHECK_EQ(this->max_rep_level(), 0); if (!this->has_values_to_process() || num_records == 0) return 0; @@ -2016,9 +1986,9 @@ TypedRecordReader::SkipRecordsInBufferNonR return skipped_records; } -template +template int64_t -TypedRecordReader::DelimitAndSkipRecordsInBuffer( +TypedRecordReader::DelimitAndSkipRecordsInBuffer( int64_t num_records) { if (num_records == 0) return 0; // Look at the buffered levels, delimit them based on @@ -2038,8 +2008,8 @@ TypedRecordReader::DelimitAndSkipRecordsIn return skipped_records; } -template -int64_t TypedRecordReader::SkipRecordsRepeated( +template +int64_t TypedRecordReader::SkipRecordsRepeated( int64_t num_records) { ARROW_DCHECK_GT(this->max_rep_level(), 0); int64_t skipped_records = 0; @@ -2106,8 +2076,8 @@ int64_t TypedRecordReader::SkipRecordsRepe return skipped_records; } -template -void TypedRecordReader::ReadAndThrowAwayValues( +template +void TypedRecordReader::ReadAndThrowAwayValues( int64_t num_values) { const int64_t values_read = this->current_decoder_.Skip(num_values); if (values_read < num_values) { @@ -2117,8 +2087,8 @@ void TypedRecordReader::ReadAndThrowAwayVa } } -template -int64_t TypedRecordReader::SkipRecords( +template +int64_t TypedRecordReader::SkipRecords( int64_t num_records) { if (num_records == 0) return 0; @@ -2146,9 +2116,9 @@ int64_t TypedRecordReader::SkipRecords( return skipped_records; } -template +template std::shared_ptr -TypedRecordReader::ReleaseIsValid() { +TypedRecordReader::ReleaseIsValid() { if (nullable_values()) { const auto bit_count = bit_util::BytesForBits(values_written()); auto result = valid_bits_; @@ -2161,8 +2131,8 @@ TypedRecordReader::ReleaseIsValid() { } } -template -int64_t TypedRecordReader::DelimitRecords( +template +int64_t TypedRecordReader::DelimitRecords( int64_t num_records, int64_t* values_seen) { if (ARROW_PREDICT_FALSE(num_records == 0 || levels_position_ == levels_written_)) { *values_seen = 0; @@ -2221,14 +2191,14 @@ int64_t TypedRecordReader::DelimitRecords( return records_read; } -template -void TypedRecordReader::Reserve(int64_t capacity) { +template +void TypedRecordReader::Reserve(int64_t capacity) { ReserveLevels(capacity); ReserveValuesAndIsValid(capacity); } -template -void TypedRecordReader::ReserveLevels( +template +void TypedRecordReader::ReserveLevels( int64_t extra_levels) { if (this->max_def_level() > 0) { const int64_t new_levels_capacity = @@ -2250,8 +2220,8 @@ void TypedRecordReader::ReserveLevels( } } -template -void TypedRecordReader::ReserveValuesAndIsValid( +template +void TypedRecordReader::ReserveValuesAndIsValid( int64_t extra_values) { values_.ReserveValues(extra_values); if (nullable_values() && !read_dense_for_nullable_) { @@ -2267,8 +2237,8 @@ void TypedRecordReader::ReserveValuesAndIs } } -template -void TypedRecordReader::Reset() { +template +void TypedRecordReader::Reset() { ResetValues(); if (levels_written_ > 0) { @@ -2279,16 +2249,16 @@ void TypedRecordReader::Reset() { // Call Finish on the binary builders to reset them } -template -void TypedRecordReader::SetPageReader( +template +void TypedRecordReader::SetPageReader( std::unique_ptr reader) { at_record_start_ = true; this->pager_ = std::move(reader); ResetDecoders(); } -template -void TypedRecordReader::ReadValuesSpaced( +template +void TypedRecordReader::ReadValuesSpaced( int64_t values_with_nulls, int64_t null_count) { uint8_t* valid_bits = valid_bits_->mutable_data(); const int64_t valid_bits_offset = values_written(); @@ -2299,16 +2269,16 @@ void TypedRecordReader::ReadValuesSpaced( CheckNumberDecoded(num_decoded, values_with_nulls); } -template -void TypedRecordReader::ReadValuesDense( +template +void TypedRecordReader::ReadValuesDense( int64_t values_to_read) { int64_t num_decoded = this->current_decoder_->Decode(values_.write_start(), static_cast(values_to_read)); CheckNumberDecoded(num_decoded, values_to_read); } -template -int64_t TypedRecordReader::ReadRepeatedRecords( +template +int64_t TypedRecordReader::ReadRepeatedRecords( int64_t num_records, int64_t* values_to_read, int64_t* null_count) { const int64_t start_levels_position = levels_position_; // Note that repeated records may be required or nullable. If they have @@ -2329,8 +2299,8 @@ int64_t TypedRecordReader::ReadRepeatedRec return records_read; } -template -int64_t TypedRecordReader::ReadOptionalRecords( +template +int64_t TypedRecordReader::ReadOptionalRecords( int64_t num_records, int64_t* values_to_read, int64_t* null_count) { const int64_t start_levels_position = levels_position_; // No repetition levels, skip delimiting logic. Each level represents a @@ -2352,8 +2322,8 @@ int64_t TypedRecordReader::ReadOptionalRec return records_read; } -template -void TypedRecordReader::ReadDenseForOptional( +template +void TypedRecordReader::ReadDenseForOptional( int64_t start_levels_position, int64_t* values_to_read) { // levels_position_ must already be incremented based on number of records // read. @@ -2366,8 +2336,8 @@ void TypedRecordReader::ReadDenseForOption ReadValuesDense(*values_to_read); } -template -void TypedRecordReader:: +template +void TypedRecordReader:: ReadSpacedForOptionalOrRepeated(int64_t start_levels_position, int64_t* values_to_read, int64_t* null_count) { // levels_position_ must already be incremented based on number of records @@ -2387,8 +2357,8 @@ void TypedRecordReader:: ReadValuesSpaced(validity_io.values_read, *null_count); } -template -int64_t TypedRecordReader::ReadRecordData( +template +int64_t TypedRecordReader::ReadRecordData( int64_t num_records) { // Conservative upper bound const int64_t possible_num_values = @@ -2442,8 +2412,8 @@ int64_t TypedRecordReader::ReadRecordData( return records_read; } -template -void TypedRecordReader::DebugPrintState() { +template +void TypedRecordReader::DebugPrintState() { const int16_t* def_levels = this->def_levels(); const int16_t* rep_levels = this->rep_levels(); const int64_t total_levels_read = levels_position_; @@ -2478,8 +2448,8 @@ void TypedRecordReader::DebugPrintState() std::cout << std::endl; } -template -void TypedRecordReader::ResetValues() { +template +void TypedRecordReader::ResetValues() { if (values_written() > 0) { values_.ResetValues(); PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); @@ -2491,8 +2461,9 @@ void TypedRecordReader::ResetValues() { * RequiredTypedRecordReader * *******************************/ -template , +// TODO devirtualized ReadDense calls via CRTP + del Hooks + if constexpr +// TODO can we reduce some code share with TypedRecordREader ? +template , bool kReadDictionary = false> class RequiredTypedRecordReader : public ColumnReaderImplBase, virtual public RecordReader, @@ -2571,16 +2542,15 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase -const void* -RequiredTypedRecordReader::ReadDictionary( +template +const void* RequiredTypedRecordReader::ReadDictionary( int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; @@ -2601,8 +2571,8 @@ RequiredTypedRecordReader::ReadDictionary( return reinterpret_cast(dictionary); } -template -int64_t RequiredTypedRecordReader::ReadRecords( +template +int64_t RequiredTypedRecordReader::ReadRecords( int64_t num_records) { if (num_records <= 0) { return 0; @@ -2632,30 +2602,30 @@ int64_t RequiredTypedRecordReader::ReadRec return records_read; } -template -void RequiredTypedRecordReader::ReadValuesDense( +template +void RequiredTypedRecordReader::ReadValuesDense( int64_t values_to_read) { const int64_t num_decoded = this->current_decoder_->Decode( values_.write_start(), static_cast(values_to_read)); CheckNumberDecoded(num_decoded, values_to_read); } -template -void RequiredTypedRecordReader::Reset() { +template +void RequiredTypedRecordReader::Reset() { if (values_written() > 0) { values_.ResetValues(); } } -template -void RequiredTypedRecordReader::SetPageReader( +template +void RequiredTypedRecordReader::SetPageReader( std::unique_ptr reader) { this->pager_ = std::move(reader); ResetDecoders(); } -template -void RequiredTypedRecordReader::DebugPrintState() { +template +void RequiredTypedRecordReader::DebugPrintState() { std::cout << "values: "; if constexpr (can_cout) { const T* vals = reinterpret_cast(this->values()); @@ -2670,6 +2640,37 @@ void RequiredTypedRecordReader::DebugPrint std::cout << std::endl; } +/******************** + * NoOpValuesSink * + ********************/ + +template +class NoOpValuesSink : private ValueSinkCursor { + public: + using value_type = T; + + using ValueSinkCursor::capacity; + using ValueSinkCursor::values_count; + + explicit NoOpValuesSink(MemoryPool* /* pool */) {} + + value_type* data() const { return nullptr; } + + value_type* write_start() { return nullptr; } + + void mark_values_as_written(int64_t extra_values) { + set_values_count(values_count() + extra_values); + } + + std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { + return nullptr; + } + + void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } + + void ResetValues() { ValueSinkCursor::operator=({}); } +}; + /************************ * record_reader_base * ************************/ @@ -2680,15 +2681,15 @@ struct record_reader_base; template struct record_reader_base { using c_type = typename DType::c_type; - using ValuesBuffer = ReadValuesNoBuffer; - using type = RequiredTypedRecordReader; + using ValueSink = NoOpValuesSink; + using type = RequiredTypedRecordReader; }; template struct record_reader_base { using c_type = typename DType::c_type; - using ValuesBuffer = ReadValuesNoBuffer; - using type = TypedRecordReader; + using ValueSink = NoOpValuesSink; + using type = TypedRecordReader; }; template @@ -3033,7 +3034,9 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* using FLBAReader = FLBARecordReader; return std::make_shared(descr, pool); } else { - return std::make_shared>(descr, pool); + using c_type = typename DType::c_type; + using Reader = RequiredTypedRecordReader, false>; + return std::make_shared(descr, pool); } } if constexpr (std::is_same_v) { @@ -3041,7 +3044,7 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } else { using c_type = typename DType::c_type; - using Reader = TypedRecordReader, false>; + using Reader = TypedRecordReader, false>; return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } } From f93b9bc0d9c5e677ef5529801082a6f5a0e31aec Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 17 Jul 2026 11:54:39 +0200 Subject: [PATCH 15/33] Split ReserveValuesAndIsValid --- cpp/src/parquet/column_reader.cc | 51 +++++++++++++++++++------------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 3bfb2d2b36b..8e90ea7455b 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1615,7 +1615,7 @@ class ReadValuesHooks { public: virtual ~ReadValuesHooks() = default; - virtual void ReserveValuesAndIsValid(int64_t extra_values) = 0; + virtual void ReserveValues(int64_t extra_values) = 0; /// Decode `values_to_read` non-null values virtual void ReadValuesDense(int64_t values_to_read) = 0; @@ -1730,7 +1730,9 @@ class TypedRecordReader : public ColumnReaderImplBase, void ReserveLevels(int64_t extra_levels); - void ReserveValuesAndIsValid(int64_t extra_values) override; + void ReserveValues(int64_t extra_values) override; + + void ReserveIsValid(int64_t extra_values); void Reset() override; @@ -2192,9 +2194,10 @@ int64_t TypedRecordReader::DelimitRecords( } template -void TypedRecordReader::Reserve(int64_t capacity) { - ReserveLevels(capacity); - ReserveValuesAndIsValid(capacity); +void TypedRecordReader::Reserve(int64_t extra_values) { + ReserveLevels(extra_values); + ReserveValues(extra_values); + ReserveIsValid(extra_values); } template @@ -2221,18 +2224,27 @@ void TypedRecordReader::ReserveLevels( } template -void TypedRecordReader::ReserveValuesAndIsValid( +void TypedRecordReader::ReserveValues( int64_t extra_values) { values_.ReserveValues(extra_values); - if (nullable_values() && !read_dense_for_nullable_) { - int64_t valid_bytes_new = bit_util::BytesForBits(values_.capacity()); - if (valid_bits_->size() < valid_bytes_new) { - int64_t valid_bytes_old = bit_util::BytesForBits(values_written()); - PARQUET_THROW_NOT_OK(valid_bits_->Resize(valid_bytes_new, /*shrink_to_fit=*/false)); +} +template +void TypedRecordReader::ReserveIsValid( + int64_t extra_values) { + if (nullable_values() && !read_dense_for_nullable_) { + const int64_t current_byte_capacity = valid_bits_->size(); + const int64_t bit_capacity = compute_capacity_pow2( + /* capacity= */ 8 * current_byte_capacity, + /* size= */ values_written(), + /* extra_size= */ extra_values); + const int64_t byte_capacity = bit_util::BytesForBits(bit_capacity); + if (current_byte_capacity < byte_capacity) { + PARQUET_THROW_NOT_OK(valid_bits_->Resize(byte_capacity, /*shrink_to_fit=*/false)); // Avoid valgrind warnings - memset(valid_bits_->mutable_data() + valid_bytes_old, 0, - static_cast(valid_bytes_new - valid_bytes_old)); + const int64_t old_valid_bytes = bit_util::BytesForBits(values_written()); + std::memset(valid_bits_->mutable_data() + old_valid_bytes, 0, + static_cast(byte_capacity - old_valid_bytes)); } } } @@ -2363,7 +2375,8 @@ int64_t TypedRecordReader::ReadRecordData( // Conservative upper bound const int64_t possible_num_values = std::max(num_records, levels_written_ - levels_position_); - ReserveValuesAndIsValid(static_cast(possible_num_values)); + ReserveValues(possible_num_values); + ReserveIsValid(possible_num_values); const int64_t start_levels_position = levels_position_; @@ -2511,7 +2524,7 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase ReleaseIsValid() final { return nullptr; } - void Reserve(int64_t extra_values) final { ReserveValuesAndIsValid(extra_values); } + void Reserve(int64_t extra_values) final { ReserveValues(extra_values); } void Reset() final; @@ -2530,7 +2543,7 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase, } protected: - void ReserveValuesAndIsValid(int64_t extra_values) override { - Base::ReserveValuesAndIsValid(extra_values); + void ReserveValues(int64_t extra_values) override { PARQUET_THROW_NOT_OK(array_builder_.Reserve(extra_values)); } @@ -2853,8 +2865,7 @@ class ByteArrayChunkedRecordReader final } protected: - void ReserveValuesAndIsValid(int64_t extra_values) override { - Base::ReserveValuesAndIsValid(extra_values); + void ReserveValues(int64_t extra_values) override { PARQUET_THROW_NOT_OK(accumulator_.builder->Reserve(extra_values)); } From ff3e0af8e87077ca12278dfe36fc1abc39ffd288 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 09:48:12 +0200 Subject: [PATCH 16/33] Devirtualize hooks with CRTP --- cpp/src/parquet/column_reader.cc | 567 ++++++++++++++++--------------- 1 file changed, 293 insertions(+), 274 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 8e90ea7455b..b63a0d37aaf 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -784,8 +784,10 @@ class ValueSinkBuffer : private ValueSinkCursor { value_type* data() const { return values_->mutable_data_as(); } + // TODO can we remove that? value_type* write_start() { return data() + values_count(); } + // TODO can we remove that? void mark_values_as_written(int64_t extra_values) { set_values_count(values_count() + extra_values); } @@ -795,6 +797,17 @@ class ValueSinkBuffer : private ValueSinkCursor { set_values_count(values_count() - count); } + [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { + return decoder.Decode(write_start(), batch_size); + } + + [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, + int32_t null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + return decoder.DecodeSpaced(write_start(), batch_size, null_count, valid_bits, + valid_bits_offset); + } + std::shared_ptr ReleaseValues(MemoryPool* pool) { // TODO should we set values_written to zero? auto result = values_; @@ -1601,47 +1614,23 @@ concept can_cout = requires(std::ostream& os, const T& value) { os << value; }; // NOLINT(readability/braces) -/********************* - * ReadValuesHooks * - *********************/ - -/// Hooks a record reader calls to materialize decoded values. -/// -/// The default implementations in `TypedRecordReader` and -/// `RequiredTypedRecordReader` decode into the `values_` buffer. Leaf readers -/// for BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY columns override them to decode -/// directly into Arrow builders instead. -class ReadValuesHooks { - public: - virtual ~ReadValuesHooks() = default; - - virtual void ReserveValues(int64_t extra_values) = 0; - - /// Decode `values_to_read` non-null values - virtual void ReadValuesDense(int64_t values_to_read) = 0; - - /// Decode a batch of `values_with_nulls` value slots, `null_count` - /// of which are null - virtual void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) = 0; -}; - /*********************** * TypedRecordReader * ***********************/ template class TypedRecordReader : public ColumnReaderImplBase, - virtual public RecordReader, - public ReadValuesHooks { + virtual public RecordReader { public: using T = typename DType::c_type; using Base = ColumnReaderImplBase; + TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, - bool read_dense_for_nullable) + bool read_dense_for_nullable, ValueSink value_sink) : Base(descr, pool, LevelDecoder(descr->max_definition_level()), LevelDecoder(descr->max_repetition_level())), valid_bits_(AllocateBuffer(pool)), - values_(pool), + value_sink_(std::move(value_sink)), leaf_info_(leaf_info), def_levels_(AllocateBuffer(pool)), rep_levels_(AllocateBuffer(pool)), @@ -1670,9 +1659,9 @@ class TypedRecordReader : public ColumnReaderImplBase, bool read_dense_for_nullable() const final { return read_dense_for_nullable_; } - uint8_t* values() const final { return reinterpret_cast(values_.data()); } + uint8_t* values() const final { return reinterpret_cast(value_sink_.data()); } - int64_t values_written() const final { return values_.values_count(); } + int64_t values_written() const final { return value_sink_.values_count(); } const void* ReadDictionary(int32_t* dictionary_length) override; @@ -1713,7 +1702,7 @@ class TypedRecordReader : public ColumnReaderImplBase, bool has_values_to_process() const { return levels_position_ < levels_written_; } std::shared_ptr ReleaseValues() override { - return values_.ReleaseValues(this->pool_); + return value_sink_.ReleaseValues(this->pool_); } std::shared_ptr ReleaseIsValid() override; @@ -1730,7 +1719,7 @@ class TypedRecordReader : public ColumnReaderImplBase, void ReserveLevels(int64_t extra_levels); - void ReserveValues(int64_t extra_values) override; + void ReserveValues(int64_t extra_values) { value_sink_.ReserveValues(extra_values); } void ReserveIsValid(int64_t extra_values); @@ -1745,9 +1734,9 @@ class TypedRecordReader : public ColumnReaderImplBase, // Dictionary decoders must be reset when advancing row groups void ResetDecoders() { this->decoders_.clear(); } - void ReadValuesSpaced(int64_t values_with_nulls, int64_t null_count) override; + void ReadValuesSpaced(int32_t values_with_nulls, int32_t null_count); - void ReadValuesDense(int64_t values_to_read) override; + void ReadValuesDense(int32_t values_to_read); // Reads repeated records and returns number of records read. Fills in // values_to_read and null_count. @@ -1782,8 +1771,11 @@ class TypedRecordReader : public ColumnReaderImplBase, /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; + auto value_sink() -> ValueSink& { return value_sink_; } + auto value_sink() const -> const ValueSink& { return value_sink_; } + private: - ValueSink values_; + ValueSink value_sink_; LevelInfo leaf_info_; /// \brief Buffer for definition levels. May contain more levels than @@ -2223,12 +2215,6 @@ void TypedRecordReader::ReserveLevels( } } -template -void TypedRecordReader::ReserveValues( - int64_t extra_values) { - values_.ReserveValues(extra_values); -} - template void TypedRecordReader::ReserveIsValid( int64_t extra_values) { @@ -2271,22 +2257,20 @@ void TypedRecordReader::SetPageReader( template void TypedRecordReader::ReadValuesSpaced( - int64_t values_with_nulls, int64_t null_count) { - uint8_t* valid_bits = valid_bits_->mutable_data(); - const int64_t valid_bits_offset = values_written(); - - int64_t num_decoded = this->current_decoder_->DecodeSpaced( - values_.write_start(), static_cast(values_with_nulls), - static_cast(null_count), valid_bits, valid_bits_offset); - CheckNumberDecoded(num_decoded, values_with_nulls); + int32_t values_with_nulls, int32_t null_count) { + const auto decoded = value_sink_.ReadValuesSpaced( + *this->current_decoder_.get(), values_with_nulls, null_count, + /* valid_bits= */ valid_bits_->mutable_data(), /* valid_bits_offset= */ + values_written()); + CheckNumberDecoded(decoded, values_with_nulls); } template void TypedRecordReader::ReadValuesDense( - int64_t values_to_read) { - int64_t num_decoded = this->current_decoder_->Decode(values_.write_start(), - static_cast(values_to_read)); - CheckNumberDecoded(num_decoded, values_to_read); + int32_t batch_size) { + const auto decoded = + value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size); + CheckNumberDecoded(decoded, batch_size); } template @@ -2302,7 +2286,8 @@ int64_t TypedRecordReader::ReadRepeatedRecord // have associated levels. int64_t records_read = DelimitRecords(num_records, values_to_read); if (!nullable_values() || read_dense_for_nullable_) { - ReadValuesDense(*values_to_read); + // This is only reading in the current page so this fits in an int32. + ReadValuesDense(clamp_to(*values_to_read)); // null_count is always 0 for required. ARROW_DCHECK_EQ(*null_count, 0); } else { @@ -2345,7 +2330,8 @@ void TypedRecordReader::ReadDenseForOptional( 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()); - ReadValuesDense(*values_to_read); + // This is only reading in the current page so this fits in an int32. + ReadValuesDense(clamp_to(*values_to_read)); } template @@ -2366,7 +2352,9 @@ void TypedRecordReader:: *null_count = validity_io.null_count; ARROW_DCHECK_GE(*values_to_read, 0); ARROW_DCHECK_GE(*null_count, 0); - ReadValuesSpaced(validity_io.values_read, *null_count); + // This is only reading in the current page so this fits in an int32. + ReadValuesSpaced(clamp_to(validity_io.values_read), + clamp_to(*null_count)); } template @@ -2397,7 +2385,8 @@ int64_t TypedRecordReader::ReadRecordData( } else { ARROW_DCHECK(!nullable_values()); values_to_read = num_records; - ReadValuesDense(values_to_read); + // This is only reading in the current page so this fits in an int32. + ReadValuesDense(clamp_to(values_to_read)); records_read = num_records; // We don't need to update null_count, since it is 0. } @@ -2407,10 +2396,10 @@ int64_t TypedRecordReader::ReadRecordData( ARROW_DCHECK_GE(null_count, 0); if (read_dense_for_nullable_) { - values_.mark_values_as_written(values_to_read); + value_sink_.mark_values_as_written(values_to_read); ARROW_DCHECK_EQ(null_count, 0); } else { - values_.mark_values_as_written(values_to_read + null_count); + value_sink_.mark_values_as_written(values_to_read + null_count); null_count_ += null_count; } // Total values, including null spaces, if any @@ -2464,7 +2453,7 @@ void TypedRecordReader::DebugPrintState() { template void TypedRecordReader::ResetValues() { if (values_written() > 0) { - values_.ResetValues(); + value_sink_.ResetValues(); PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); null_count_ = 0; } @@ -2474,27 +2463,28 @@ void TypedRecordReader::ResetValues() { * RequiredTypedRecordReader * *******************************/ -// TODO devirtualized ReadDense calls via CRTP + del Hooks + if constexpr // TODO can we reduce some code share with TypedRecordREader ? template , bool kReadDictionary = false> class RequiredTypedRecordReader : public ColumnReaderImplBase, - virtual public RecordReader, - public ReadValuesHooks { + virtual public RecordReader { public: using T = typename DType::c_type; using Base = ColumnReaderImplBase; - RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool) + RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool, + ValueSink value_sink) : Base(descr, pool, LevelDecoder(descr->max_definition_level()), LevelDecoder(descr->max_repetition_level())), - values_(pool) { + value_sink_(std::move(value_sink)) { RequiredTypedRecordReader::Reset(); + ARROW_DCHECK_EQ(descr->max_definition_level(), 0); + ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); } - uint8_t* values() const final { return reinterpret_cast(values_.data()); } + uint8_t* values() const final { return reinterpret_cast(value_sink_.data()); } - int64_t values_written() const final { return values_.values_count(); } + int64_t values_written() const final { return value_sink_.values_count(); } int16_t* def_levels() const final { return nullptr; } @@ -2519,7 +2509,7 @@ class RequiredTypedRecordReader : public ColumnReaderImplBaseSkip(num_records); } std::shared_ptr ReleaseValues() final { - return values_.ReleaseValues(this->pool_); + return value_sink_.ReleaseValues(this->pool_); } std::shared_ptr ReleaseIsValid() final { return nullptr; } @@ -2540,22 +2530,15 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase ValueSink& { return value_sink_; } + auto value_sink() const -> const ValueSink& { return value_sink_; } - void ReadValuesDense(int64_t values_to_read) override; + private: + ValueSink value_sink_; - // A required column never contains nulls, so values are never read spaced. - void ReadValuesSpaced(int64_t /*values_with_nulls*/, int64_t /*null_count*/) override { - throw ParquetException("Spaced values read in required column"); - } + void ReserveValues(int64_t extra_values) { value_sink_.ReserveValues(extra_values); } - private: - ValueSink values_; + void ReadValuesDense(int32_t values_to_read); }; /********************************************** @@ -2606,7 +2589,7 @@ int64_t RequiredTypedRecordReader::ReadRecord this->available_values_current_page()); ReadValuesDense(batch_size); - values_.mark_values_as_written(batch_size); + value_sink_.mark_values_as_written(batch_size); this->ConsumeBufferedValues(batch_size); records_read += batch_size; @@ -2617,16 +2600,16 @@ int64_t RequiredTypedRecordReader::ReadRecord template void RequiredTypedRecordReader::ReadValuesDense( - int64_t values_to_read) { - const int64_t num_decoded = this->current_decoder_->Decode( - values_.write_start(), static_cast(values_to_read)); - CheckNumberDecoded(num_decoded, values_to_read); + int32_t batch_size) { + const auto decoded = + value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size); + CheckNumberDecoded(decoded, batch_size); } template void RequiredTypedRecordReader::Reset() { if (values_written() > 0) { - values_.ResetValues(); + value_sink_.ResetValues(); } } @@ -2653,19 +2636,19 @@ void RequiredTypedRecordReader::DebugPrintSta std::cout << std::endl; } -/******************** - * NoOpValuesSink * - ********************/ +/********************* + * ArrayValuesSink * + *********************/ -template -class NoOpValuesSink : private ValueSinkCursor { +template +class ArrayValuesSink : private ValueSinkCursor { public: using value_type = T; using ValueSinkCursor::capacity; using ValueSinkCursor::values_count; - explicit NoOpValuesSink(MemoryPool* /* pool */) {} + explicit ArrayValuesSink(BuilderType builder) : builder_{std::move(builder)} {} value_type* data() const { return nullptr; } @@ -2679,40 +2662,62 @@ class NoOpValuesSink : private ValueSinkCursor { return nullptr; } - void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } + void ReserveValues(int64_t extra_values) { + fit_capacity_for_extra(extra_values); + // The chunked accumulator is not an ArrayBuilder and cannot be reserved. + if constexpr (requires { builder_.Reserve(extra_values); }) { + PARQUET_THROW_NOT_OK(builder_.Reserve(extra_values)); + } + } void ResetValues() { ValueSinkCursor::operator=({}); } -}; -/************************ - * record_reader_base * - ************************/ + [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { + // TODO once we have a validity sink: we can reset the validity to save some space + return decoder.DecodeArrowNonNull(batch_size, &builder_); + } -template -struct record_reader_base; + [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, + int32_t null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + // TODO once we have a validity sink: we can reset the validity to save some space + const int64_t decoded = decoder.DecodeArrow(batch_size, null_count, valid_bits, + valid_bits_offset, &builder_); + // `DecodeArrow` only counts the non-null values, but the caller expects the + // number of values written, nulls included. + ARROW_DCHECK_EQ(decoded, batch_size - null_count); + return decoded + null_count; + } -template -struct record_reader_base { - using c_type = typename DType::c_type; - using ValueSink = NoOpValuesSink; - using type = RequiredTypedRecordReader; -}; + auto builder() -> BuilderType& { return builder_; } + auto builder() const -> const BuilderType& { return builder_; } -template -struct record_reader_base { - using c_type = typename DType::c_type; - using ValueSink = NoOpValuesSink; - using type = TypedRecordReader; + private: + BuilderType builder_; }; -template -using record_reader_base_t = - typename record_reader_base::type; - /********************** * FLBARecordReader * **********************/ +template +using record_reader_base_t = + std::conditional_t, + TypedRecordReader>; + +template +struct flba_record_reader_base { + using DType = FLBAType; + using c_type = typename DType::c_type; + using Builder = ::arrow::FixedSizeBinaryBuilder; + using ValueSink = ArrayValuesSink; + using type = record_reader_base_t; +}; + +template +using flba_record_reader_base_t = typename flba_record_reader_base::type; + /// Reads fixed length byte array values into a FixedSizeBinaryBuilder. /// /// `kRequired` selects the base class: RequiredTypedRecordReader for @@ -2724,68 +2729,38 @@ using record_reader_base_t = /// values written. The `valid_bits_` buffer, if any, is consumed by each /// spaced decode. template -class FLBARecordReader final : public record_reader_base_t, +class FLBARecordReader final : public flba_record_reader_base_t, virtual public BinaryRecordReader { public: - using Base = record_reader_base_t; + using Base = flba_record_reader_base_t; FLBARecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) requires(!kRequired) - : Base(descr, leaf_info, pool, read_dense_for_nullable), - byte_width_(descr->type_length()), - type_(::arrow::fixed_size_binary(byte_width_)), - array_builder_(type_, pool) { + : Base(descr, leaf_info, pool, read_dense_for_nullable, MakeSink(descr, pool)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); } FLBARecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) requires(kRequired) - : Base(descr, pool), - byte_width_(descr->type_length()), - type_(::arrow::fixed_size_binary(byte_width_)), - array_builder_(type_, pool) { + : Base(descr, pool, MakeSink(descr, pool)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); - ARROW_DCHECK_EQ(descr->max_definition_level(), 0); - ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); } ::arrow::ArrayVector GetBuilderChunks() override { - PARQUET_ASSIGN_OR_THROW(auto chunk, array_builder_.Finish()); + PARQUET_ASSIGN_OR_THROW(auto chunk, builder().Finish()); return ::arrow::ArrayVector{std::move(chunk)}; } - protected: - void ReserveValues(int64_t extra_values) override { - PARQUET_THROW_NOT_OK(array_builder_.Reserve(extra_values)); - } - - void ReadValuesDense(int64_t values_to_read) override { - int64_t num_decoded = this->current_decoder_->DecodeArrowNonNull( - static_cast(values_to_read), &array_builder_); - CheckNumberDecoded(num_decoded, values_to_read); - if constexpr (!kRequired) { - this->ResetValues(); - } - } + private: + using Builder = typename flba_record_reader_base::Builder; + using ValueSink = typename flba_record_reader_base::ValueSink; - void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - if constexpr (kRequired) { - // A required column never contains nulls: the base implementation throws. - Base::ReadValuesSpaced(values_to_read, null_count); - } else { - int64_t num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - this->valid_bits_->mutable_data(), this->values_written(), &array_builder_); - CheckNumberDecoded(num_decoded, values_to_read - null_count); - this->ResetValues(); - } + static auto MakeSink(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) { + return ValueSink(Builder(::arrow::fixed_size_binary(descr->type_length()), pool)); } - private: - const int byte_width_; - std::shared_ptr<::arrow::DataType> type_; - ::arrow::FixedSizeBinaryBuilder array_builder_; + auto builder() -> Builder& { return this->value_sink().builder(); } }; /*************************** @@ -2794,8 +2769,8 @@ class FLBARecordReader final : public record_reader_base_t, /// Create the Arrow builder for reading a Parquet BYTE_ARRAY column as the /// given Arrow type (BINARY by default). -std::unique_ptr<::arrow::ArrayBuilder> MakeByteArrayBuilder( - const std::shared_ptr<::arrow::DataType>& arrow_type, ::arrow::MemoryPool* pool) { +std::unique_ptr<::arrow::ArrayBuilder> MakeByteArrayBuilder(::arrow::DataType* arrow_type, + ::arrow::MemoryPool* pool) { auto arrow_binary_type = arrow_type ? arrow_type->id() : ::arrow::Type::BINARY; switch (arrow_binary_type) { case ::arrow::Type::BINARY: @@ -2816,6 +2791,19 @@ std::unique_ptr<::arrow::ArrayBuilder> MakeByteArrayBuilder( } } +template +struct byte_array_chunked_record_reader { + using DType = ByteArrayType; + using c_type = typename DType::c_type; + using Builder = typename EncodingTraits::Accumulator; + using ValueSink = ArrayValuesSink; + using type = record_reader_base_t; +}; + +template +using byte_array_chunked_record_reader_t = + typename byte_array_chunked_record_reader::type; + /// Reads variable length byte array values into a chunked binary builder. /// /// `kRequired` selects the base class: RequiredTypedRecordReader for @@ -2829,181 +2817,209 @@ std::unique_ptr<::arrow::ArrayBuilder> MakeByteArrayBuilder( /// is used to store the values. template class ByteArrayChunkedRecordReader final - : public record_reader_base_t, + : public byte_array_chunked_record_reader_t, virtual public BinaryRecordReader { public: - using Base = record_reader_base_t; + using Base = byte_array_chunked_record_reader_t; ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable, const std::shared_ptr<::arrow::DataType>& arrow_type) requires(!kRequired) - : Base(descr, leaf_info, pool, read_dense_for_nullable) { + : Base(descr, leaf_info, pool, read_dense_for_nullable, + MakeSink(pool, arrow_type)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); - accumulator_.builder = MakeByteArrayBuilder(arrow_type, pool); } ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, const std::shared_ptr<::arrow::DataType>& arrow_type) requires(kRequired) - : Base(descr, pool) { + : Base(descr, pool, MakeSink(pool, arrow_type)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); - ARROW_DCHECK_EQ(descr->max_definition_level(), 0); - ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); - accumulator_.builder = MakeByteArrayBuilder(arrow_type, pool); } ::arrow::ArrayVector GetBuilderChunks() override { - ::arrow::ArrayVector result = accumulator_.chunks; - if (result.empty() || accumulator_.builder->length() > 0) { + ::arrow::ArrayVector result = accumulator().chunks; + if (result.empty() || accumulator().builder->length() > 0) { std::shared_ptr<::arrow::Array> last_chunk; - PARQUET_THROW_NOT_OK(accumulator_.builder->Finish(&last_chunk)); + PARQUET_THROW_NOT_OK(accumulator().builder->Finish(&last_chunk)); result.push_back(std::move(last_chunk)); } - accumulator_.chunks = {}; + accumulator().chunks = {}; return result; } - protected: - void ReserveValues(int64_t extra_values) override { - PARQUET_THROW_NOT_OK(accumulator_.builder->Reserve(extra_values)); + private: + using Builder = typename byte_array_chunked_record_reader::Builder; + using ValueSink = typename byte_array_chunked_record_reader::ValueSink; + + static auto MakeSink(::arrow::MemoryPool* pool, + const std::shared_ptr<::arrow::DataType>& arrow_type) { + Builder accumulator = {}; + accumulator.builder = MakeByteArrayBuilder(arrow_type.get(), pool); + return ValueSink(std::move(accumulator)); } - void ReadValuesDense(int64_t values_to_read) override { - int64_t num_decoded = this->current_decoder_->DecodeArrowNonNull( - static_cast(values_to_read), &accumulator_); - CheckNumberDecoded(num_decoded, values_to_read); - if constexpr (!kRequired) { - this->ResetValues(); - } + auto accumulator() -> Builder& { return this->value_sink().builder(); } +}; + +/// Decodes byte array values into a ::arrow::BinaryDictionary32Builder. +/// +/// If the current decoder is dictionary encoded, the dictionary indices are +/// decoded directly, otherwise the values are decoded and looked up in the +/// builder's memo table. +class ByteArrayDictionaryValuesSink : private ValueSinkCursor { + public: + using value_type = ByteArray; + using Builder = ::arrow::BinaryDictionary32Builder; + + using ValueSinkCursor::capacity; + using ValueSinkCursor::values_count; + + explicit ByteArrayDictionaryValuesSink(::arrow::MemoryPool* pool) + : builder_(std::make_unique(pool)) {} + + value_type* data() const { return nullptr; } + + void mark_values_as_written(int64_t extra_values) { + set_values_count(values_count() + extra_values); } - void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - if constexpr (kRequired) { - // A required column never contains nulls: the base implementation throws. - Base::ReadValuesSpaced(values_to_read, null_count); - } else { - int64_t num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - this->valid_bits_->mutable_data(), this->values_written(), &accumulator_); - CheckNumberDecoded(num_decoded, values_to_read - null_count); - this->ResetValues(); + std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { + return nullptr; + } + + void ReserveValues(int64_t extra_values) { fit_capacity_for_extra(extra_values); } + + void ResetValues() { ValueSinkCursor::operator=({}); } + + [[nodiscard]] int64_t ReadValuesDense(auto& decoder, int32_t batch_size) { + // TODO once we have a validity sink: we can reset the validity to save some space + if (auto* dict_decoder = dynamic_cast(&decoder)) { + MaybeWriteNewDictionary(*dict_decoder); + return dict_decoder->DecodeIndices(batch_size, builder_.get()); } + return decoder.DecodeArrowNonNull(batch_size, builder_.get()); + } + + [[nodiscard]] int64_t ReadValuesSpaced(auto& decoder, int32_t batch_size, + int32_t null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + // TODO once we have a validity sink: we can reset the validity to save some space + const int64_t decoded = [&] { + if (auto* dict_decoder = dynamic_cast(&decoder)) { + MaybeWriteNewDictionary(*dict_decoder); + return dict_decoder->DecodeIndicesSpaced(batch_size, null_count, valid_bits, + valid_bits_offset, builder_.get()); + } + return decoder.DecodeArrow(batch_size, null_count, valid_bits, valid_bits_offset, + builder_.get()); + }(); + // The decoders only count the non-null values, but the caller expects the + // number of values written, nulls included. + ARROW_DCHECK_EQ(decoded, batch_size - null_count); + return decoded + null_count; + } + + /// Bind the reader's flag signalling that a new dictionary page was read. + /// + /// The flag belongs to the reader base class, which is only constructed + /// once this sink has been handed over to it, hence the late binding. + void BindNewDictionaryFlag(bool* new_dictionary) { new_dictionary_ = new_dictionary; } + + /// Finish the builder and return all the chunks accumulated so far. + std::shared_ptr<::arrow::ChunkedArray> FlushChunks() { + FlushBuilder(); + return std::make_shared<::arrow::ChunkedArray>(std::exchange(result_chunks_, {}), + builder_->type()); } private: - // Helper data structure for accumulating builder chunks - typename EncodingTraits::Accumulator accumulator_; + using BinaryDictDecoder = DictDecoder; + + std::unique_ptr builder_; + std::vector> result_chunks_; + bool* new_dictionary_ = nullptr; + + void FlushBuilder() { + if (builder_->length() > 0) { + std::shared_ptr<::arrow::Array> chunk; + PARQUET_THROW_NOT_OK(builder_->Finish(&chunk)); + result_chunks_.emplace_back(std::move(chunk)); + + // Also clears the dictionary memo table + builder_->Reset(); + } + } + + void MaybeWriteNewDictionary(BinaryDictDecoder& decoder) { + ARROW_DCHECK_NE(new_dictionary_, nullptr); + if (*new_dictionary_) { + /// If there is a new dictionary, we may need to flush the builder, then + /// insert the new dictionary values + FlushBuilder(); + builder_->ResetFull(); + decoder.InsertDictionary(builder_.get()); + *new_dictionary_ = false; + } + } +}; + +template +struct byte_array_dictionary_record_reader { + using DType = ByteArrayType; + using ValueSink = ByteArrayDictionaryValuesSink; + using type = + record_reader_base_t; }; +template +using byte_array_dictionary_record_reader_t = + typename byte_array_dictionary_record_reader::type; + /// Reads byte array values into ::arrow::dictionary(index: int32, values: binary). /// /// `kRequired` selects the base class: RequiredTypedRecordReader for /// required (non-nullable, non-repeated) columns, TypedRecordReader /// otherwise. /// -/// If the underlying column is dictionary encoded, it will call -/// `DecodeIndices` to read, otherwise it will call `DecodeArrowNonNull` to -/// read. +/// The `values_` buffers are never used, the values are stored in the +/// dictionary builder held by the value sink. template class ByteArrayDictionaryRecordReader final - : public record_reader_base_t, + : public byte_array_dictionary_record_reader_t, virtual public DictionaryRecordReader { public: - using Base = record_reader_base_t; + using Base = byte_array_dictionary_record_reader_t; ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dense_for_nullable) requires(!kRequired) - : Base(descr, leaf_info, pool, read_dense_for_nullable), builder_(pool) { + : Base(descr, leaf_info, pool, read_dense_for_nullable, ValueSink(pool)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); + BindNewDictionaryFlag(); } ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) requires(kRequired) - : Base(descr, pool), builder_(pool) { + : Base(descr, pool, ValueSink(pool)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); ARROW_DCHECK_EQ(descr->max_definition_level(), 0); ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); + BindNewDictionaryFlag(); } std::shared_ptr<::arrow::ChunkedArray> GetResult() override { - FlushBuilder(); - std::vector> result; - std::swap(result, result_chunks_); - return std::make_shared<::arrow::ChunkedArray>(std::move(result), builder_.type()); - } - - protected: - void FlushBuilder() { - if (builder_.length() > 0) { - std::shared_ptr<::arrow::Array> chunk; - PARQUET_THROW_NOT_OK(builder_.Finish(&chunk)); - result_chunks_.emplace_back(std::move(chunk)); - - // Also clears the dictionary memo table - builder_.Reset(); - } - } - - void MaybeWriteNewDictionary() { - if (this->new_dictionary_) { - /// If there is a new dictionary, we may need to flush the builder, then - /// insert the new dictionary values - FlushBuilder(); - builder_.ResetFull(); - auto decoder = dynamic_cast(this->current_decoder_.get()); - decoder->InsertDictionary(&builder_); - this->new_dictionary_ = false; - } - } - - void ReadValuesDense(int64_t values_to_read) override { - int64_t num_decoded = 0; - if (this->current_encoding_ == Encoding::RLE_DICTIONARY) { - MaybeWriteNewDictionary(); - 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( - static_cast(values_to_read), &builder_); - } - if constexpr (!kRequired) { - // Flush values since they have been copied into the builder - this->ResetValues(); - } - CheckNumberDecoded(num_decoded, values_to_read); - } - - void ReadValuesSpaced(int64_t values_to_read, int64_t null_count) override { - if constexpr (kRequired) { - // A required column never contains nulls: the base implementation throws. - Base::ReadValuesSpaced(values_to_read, null_count); - } else { - int64_t num_decoded = 0; - if (this->current_encoding_ == Encoding::RLE_DICTIONARY) { - MaybeWriteNewDictionary(); - auto decoder = dynamic_cast(this->current_decoder_.get()); - num_decoded = decoder->DecodeIndicesSpaced( - static_cast(values_to_read), static_cast(null_count), - this->valid_bits_->mutable_data(), this->values_written(), &builder_); - } else { - num_decoded = this->current_decoder_->DecodeArrow( - static_cast(values_to_read), static_cast(null_count), - this->valid_bits_->mutable_data(), this->values_written(), &builder_); - } - ARROW_DCHECK_EQ(num_decoded, values_to_read - null_count); - // Flush values since they have been copied into the builder - this->ResetValues(); - } + return this->value_sink().FlushChunks(); } private: - using BinaryDictDecoder = DictDecoder; + using ValueSink = typename byte_array_dictionary_record_reader::ValueSink; - ::arrow::BinaryDictionary32Builder builder_; - std::vector> result_chunks_; + void BindNewDictionaryFlag() { + this->value_sink().BindNewDictionaryFlag(&this->new_dictionary_); + } }; std::shared_ptr MakeByteArrayRecordReader( @@ -3046,8 +3062,9 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* return std::make_shared(descr, pool); } else { using c_type = typename DType::c_type; - using Reader = RequiredTypedRecordReader, false>; - return std::make_shared(descr, pool); + using ValueSink = ValueSinkBuffer; + using Reader = RequiredTypedRecordReader; + return std::make_shared(descr, pool, ValueSink(pool)); } } if constexpr (std::is_same_v) { @@ -3055,8 +3072,10 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } else { using c_type = typename DType::c_type; - using Reader = TypedRecordReader, false>; - return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); + using ValueSink = ValueSinkBuffer; + using Reader = TypedRecordReader; + return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable, + ValueSink(pool)); } } } // namespace From 6f5e1fb515d6bd1e59acfce4a86b8bdba412556f Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 10:00:18 +0200 Subject: [PATCH 17/33] Simplify template --- cpp/src/parquet/column_reader.cc | 155 ++++++++++++++----------------- 1 file changed, 70 insertions(+), 85 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index b63a0d37aaf..92df405646c 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1820,9 +1820,8 @@ class TypedRecordReader : public ColumnReaderImplBase, * TypedRecordReader Implementation * **************************************/ -template -const void* TypedRecordReader::ReadDictionary( - int32_t* dictionary_length) { +template +const void* TypedRecordReader::ReadDictionary(int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; return nullptr; @@ -1836,15 +1835,14 @@ const void* TypedRecordReader::ReadDictionary << EncodingToString(this->current_encoding_); throw ParquetException(ss.str()); } - auto decoder = dynamic_cast*>(this->current_decoder_.get()); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); return reinterpret_cast(dictionary); } -template -int64_t TypedRecordReader::ReadRecords( - int64_t num_records) { +template +int64_t TypedRecordReader::ReadRecords(int64_t num_records) { if (num_records == 0) return 0; // Delimit records, then read values at the end int64_t records_read = 0; @@ -1912,9 +1910,8 @@ int64_t TypedRecordReader::ReadRecords( return records_read; } -template -void TypedRecordReader::ThrowAwayLevels( - int64_t start_levels_position) { +template +void TypedRecordReader::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); @@ -1945,9 +1942,8 @@ void TypedRecordReader::ThrowAwayLevels( levels_capacity_ -= gap; } -template -int64_t -TypedRecordReader::SkipRecordsInBufferNonRepeated( +template +int64_t TypedRecordReader::SkipRecordsInBufferNonRepeated( int64_t num_records) { ARROW_DCHECK_EQ(this->max_rep_level(), 0); if (!this->has_values_to_process() || num_records == 0) return 0; @@ -1980,9 +1976,8 @@ TypedRecordReader::SkipRecordsInBufferNonRepe return skipped_records; } -template -int64_t -TypedRecordReader::DelimitAndSkipRecordsInBuffer( +template +int64_t TypedRecordReader::DelimitAndSkipRecordsInBuffer( int64_t num_records) { if (num_records == 0) return 0; // Look at the buffered levels, delimit them based on @@ -2002,9 +1997,8 @@ TypedRecordReader::DelimitAndSkipRecordsInBuf return skipped_records; } -template -int64_t TypedRecordReader::SkipRecordsRepeated( - int64_t num_records) { +template +int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_records) { ARROW_DCHECK_GT(this->max_rep_level(), 0); int64_t skipped_records = 0; @@ -2070,9 +2064,8 @@ int64_t TypedRecordReader::SkipRecordsRepeate return skipped_records; } -template -void TypedRecordReader::ReadAndThrowAwayValues( - int64_t num_values) { +template +void TypedRecordReader::ReadAndThrowAwayValues(int64_t num_values) { const int64_t values_read = this->current_decoder_.Skip(num_values); if (values_read < num_values) { std::stringstream ss; @@ -2081,9 +2074,8 @@ void TypedRecordReader::ReadAndThrowAwayValue } } -template -int64_t TypedRecordReader::SkipRecords( - int64_t num_records) { +template +int64_t TypedRecordReader::SkipRecords(int64_t num_records) { if (num_records == 0) return 0; // Top level required field. Number of records equals to number of levels, @@ -2110,9 +2102,8 @@ int64_t TypedRecordReader::SkipRecords( return skipped_records; } -template -std::shared_ptr -TypedRecordReader::ReleaseIsValid() { +template +std::shared_ptr TypedRecordReader::ReleaseIsValid() { if (nullable_values()) { const auto bit_count = bit_util::BytesForBits(values_written()); auto result = valid_bits_; @@ -2125,9 +2116,9 @@ TypedRecordReader::ReleaseIsValid() { } } -template -int64_t TypedRecordReader::DelimitRecords( - int64_t num_records, int64_t* values_seen) { +template +int64_t TypedRecordReader::DelimitRecords(int64_t num_records, + int64_t* values_seen) { if (ARROW_PREDICT_FALSE(num_records == 0 || levels_position_ == levels_written_)) { *values_seen = 0; return 0; @@ -2185,16 +2176,15 @@ int64_t TypedRecordReader::DelimitRecords( return records_read; } -template -void TypedRecordReader::Reserve(int64_t extra_values) { +template +void TypedRecordReader::Reserve(int64_t extra_values) { ReserveLevels(extra_values); ReserveValues(extra_values); ReserveIsValid(extra_values); } -template -void TypedRecordReader::ReserveLevels( - int64_t extra_levels) { +template +void TypedRecordReader::ReserveLevels(int64_t extra_levels) { if (this->max_def_level() > 0) { const int64_t new_levels_capacity = compute_capacity_pow2(levels_capacity_, levels_written_, extra_levels); @@ -2215,9 +2205,8 @@ void TypedRecordReader::ReserveLevels( } } -template -void TypedRecordReader::ReserveIsValid( - int64_t extra_values) { +template +void TypedRecordReader::ReserveIsValid(int64_t extra_values) { if (nullable_values() && !read_dense_for_nullable_) { const int64_t current_byte_capacity = valid_bits_->size(); const int64_t bit_capacity = compute_capacity_pow2( @@ -2235,8 +2224,8 @@ void TypedRecordReader::ReserveIsValid( } } -template -void TypedRecordReader::Reset() { +template +void TypedRecordReader::Reset() { ResetValues(); if (levels_written_ > 0) { @@ -2247,17 +2236,16 @@ void TypedRecordReader::Reset() { // Call Finish on the binary builders to reset them } -template -void TypedRecordReader::SetPageReader( - std::unique_ptr reader) { +template +void TypedRecordReader::SetPageReader(std::unique_ptr reader) { at_record_start_ = true; this->pager_ = std::move(reader); ResetDecoders(); } -template -void TypedRecordReader::ReadValuesSpaced( - int32_t values_with_nulls, int32_t null_count) { +template +void TypedRecordReader::ReadValuesSpaced(int32_t values_with_nulls, + int32_t null_count) { const auto decoded = value_sink_.ReadValuesSpaced( *this->current_decoder_.get(), values_with_nulls, null_count, /* valid_bits= */ valid_bits_->mutable_data(), /* valid_bits_offset= */ @@ -2265,17 +2253,17 @@ void TypedRecordReader::ReadValuesSpaced( CheckNumberDecoded(decoded, values_with_nulls); } -template -void TypedRecordReader::ReadValuesDense( - int32_t batch_size) { +template +void TypedRecordReader::ReadValuesDense(int32_t batch_size) { const auto decoded = value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size); CheckNumberDecoded(decoded, batch_size); } -template -int64_t TypedRecordReader::ReadRepeatedRecords( - int64_t num_records, int64_t* values_to_read, int64_t* null_count) { +template +int64_t TypedRecordReader::ReadRepeatedRecords(int64_t num_records, + int64_t* values_to_read, + int64_t* null_count) { const int64_t start_levels_position = levels_position_; // Note that repeated records may be required or nullable. If they have // an optional parent in the path, they will be nullable, otherwise, @@ -2296,9 +2284,10 @@ int64_t TypedRecordReader::ReadRepeatedRecord return records_read; } -template -int64_t TypedRecordReader::ReadOptionalRecords( - int64_t num_records, int64_t* values_to_read, int64_t* null_count) { +template +int64_t TypedRecordReader::ReadOptionalRecords(int64_t num_records, + int64_t* values_to_read, + int64_t* null_count) { const int64_t start_levels_position = levels_position_; // No repetition levels, skip delimiting logic. Each level represents a // null or not null entry @@ -2319,9 +2308,9 @@ int64_t TypedRecordReader::ReadOptionalRecord return records_read; } -template -void TypedRecordReader::ReadDenseForOptional( - int64_t start_levels_position, int64_t* values_to_read) { +template +void TypedRecordReader::ReadDenseForOptional(int64_t start_levels_position, + int64_t* values_to_read) { // levels_position_ must already be incremented based on number of records // read. ARROW_DCHECK_GE(levels_position_, start_levels_position); @@ -2334,10 +2323,9 @@ void TypedRecordReader::ReadDenseForOptional( ReadValuesDense(clamp_to(*values_to_read)); } -template -void TypedRecordReader:: - ReadSpacedForOptionalOrRepeated(int64_t start_levels_position, - int64_t* values_to_read, int64_t* null_count) { +template +void TypedRecordReader::ReadSpacedForOptionalOrRepeated( + int64_t start_levels_position, int64_t* values_to_read, int64_t* null_count) { // levels_position_ must already be incremented based on number of records // read. ARROW_DCHECK_GE(levels_position_, start_levels_position); @@ -2357,9 +2345,8 @@ void TypedRecordReader:: clamp_to(*null_count)); } -template -int64_t TypedRecordReader::ReadRecordData( - int64_t num_records) { +template +int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { // Conservative upper bound const int64_t possible_num_values = std::max(num_records, levels_written_ - levels_position_); @@ -2414,8 +2401,8 @@ int64_t TypedRecordReader::ReadRecordData( return records_read; } -template -void TypedRecordReader::DebugPrintState() { +template +void TypedRecordReader::DebugPrintState() { const int16_t* def_levels = this->def_levels(); const int16_t* rep_levels = this->rep_levels(); const int64_t total_levels_read = levels_position_; @@ -2450,8 +2437,8 @@ void TypedRecordReader::DebugPrintState() { std::cout << std::endl; } -template -void TypedRecordReader::ResetValues() { +template +void TypedRecordReader::ResetValues() { if (values_written() > 0) { value_sink_.ResetValues(); PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); @@ -2545,8 +2532,8 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase -const void* RequiredTypedRecordReader::ReadDictionary( +template +const void* RequiredTypedRecordReader::ReadDictionary( int32_t* dictionary_length) { if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; @@ -2561,15 +2548,14 @@ const void* RequiredTypedRecordReader::ReadDi << EncodingToString(this->current_encoding_); throw ParquetException(ss.str()); } - auto decoder = dynamic_cast*>(this->current_decoder_.get()); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); return reinterpret_cast(dictionary); } -template -int64_t RequiredTypedRecordReader::ReadRecords( - int64_t num_records) { +template +int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records) { if (num_records <= 0) { return 0; } @@ -2598,30 +2584,29 @@ int64_t RequiredTypedRecordReader::ReadRecord return records_read; } -template -void RequiredTypedRecordReader::ReadValuesDense( - int32_t batch_size) { +template +void RequiredTypedRecordReader::ReadValuesDense(int32_t batch_size) { const auto decoded = value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size); CheckNumberDecoded(decoded, batch_size); } -template -void RequiredTypedRecordReader::Reset() { +template +void RequiredTypedRecordReader::Reset() { if (values_written() > 0) { value_sink_.ResetValues(); } } -template -void RequiredTypedRecordReader::SetPageReader( +template +void RequiredTypedRecordReader::SetPageReader( std::unique_ptr reader) { this->pager_ = std::move(reader); ResetDecoders(); } -template -void RequiredTypedRecordReader::DebugPrintState() { +template +void RequiredTypedRecordReader::DebugPrintState() { std::cout << "values: "; if constexpr (can_cout) { const T* vals = reinterpret_cast(this->values()); From 93c7b20dee1edb72b4ab34fa52b6ae996c008e66 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 10:14:01 +0200 Subject: [PATCH 18/33] Remove StackedLevelDecoder --- cpp/src/parquet/column_reader.cc | 126 ------------------------------- 1 file changed, 126 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 92df405646c..b95baaddd5a 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -846,132 +846,6 @@ class ValueSinkBuffer : private ValueSinkCursor { } }; -/************************* - * StackedLevelDecoder * - *************************/ - -// TODO do we still need this? -class StackedLevelDecoder : private LevelDecoder { - public: - using Base = LevelDecoder; - using CountUpToResult = Base::CountUpToResult; - - static constexpr int32_t kMinBatchSize = kMinLevelBatchSize; - - StackedLevelDecoder(int16_t max_level, MemoryPool* pool) - : Base{max_level}, buffer_(pool) {} - - int32_t SetData(Encoding::type encoding, int16_t max_level, int32_t num_buffered_values, - const uint8_t* data, int32_t data_size); - - void SetDataV2(int32_t num_bytes, int16_t max_level, int32_t num_buffered_values, - const uint8_t* data); - - [[nodiscard]] int32_t Decode(int32_t batch_size); - - [[nodiscard]] int32_t Skip(int32_t batch_size); - - CountUpToResult CountUpTo(int16_t value, int32_t batch_size); - - int16_t* data() const { return buffer_.data(); } - - int64_t size() const { return position_; } - - private: - ValueSinkBuffer buffer_; - int64_t position_ = 0; - - int64_t total_decoded() const { return buffer_.values_count(); } - - int32_t extra_decoded() const { - const auto count = total_decoded() - size(); - ARROW_DCHECK_LT(count, kMinBatchSize); - return static_cast(count); - } - - int16_t* extra_start() { return data() + size(); } - - void AdvanceInBuffer(int32_t batch_size) { position_ += batch_size; } - - /// Left shift the extra values in the buffer by up to the batch_size. - int32_t SkipInBuffer(int32_t batch_size); -}; - -/**************************************** - * StackedLevelDecoder Implementation * - ****************************************/ - -int32_t StackedLevelDecoder::SetData(Encoding::type encoding, int16_t max_level, - int32_t num_buffered_values, const uint8_t* data, - int32_t data_size) { - buffer_.delete_back(extra_decoded()); - return Base::SetData(encoding, max_level, num_buffered_values, data, data_size); -} - -void StackedLevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, - int32_t num_buffered_values, const uint8_t* data) { - buffer_.delete_back(extra_decoded()); - return Base::SetDataV2(num_bytes, max_level, num_buffered_values, data); -} - -int32_t StackedLevelDecoder::Decode(int32_t batch_size) { - // Advance as much as possible in the already decoded buffer. - const auto buffer_batch = std::min(batch_size, extra_decoded()); - AdvanceInBuffer(buffer_batch); - - // Check what is left to decode. - const auto decode_batch = batch_size - buffer_batch; - const auto decode_possible_batch = std::min(decode_batch, Base::remaining()); - if (decode_possible_batch == 0) { - return buffer_batch; - } - - // Decode the rest plus extra greedy values. - const auto greedy_target = std::max(kMinBatchSize, decode_batch); - const auto greedy_possible = std::min(greedy_target, Base::remaining()); - buffer_.ReserveValues(/* extra= */ greedy_possible); - const auto decoded = Base::Decode(greedy_possible, buffer_.write_start()); - buffer_.mark_values_as_written(decoded); - ARROW_DCHECK_EQ(decoded, greedy_possible); - AdvanceInBuffer(decode_possible_batch); - - return buffer_batch + decode_possible_batch; -} - -int32_t StackedLevelDecoder::SkipInBuffer(int32_t batch_size) { - const auto count = std::min(batch_size, extra_decoded()); - std::copy(extra_start() + count, extra_start() + extra_decoded(), extra_start()); - buffer_.delete_back(count); - return count; -} - -int32_t StackedLevelDecoder::Skip(int32_t batch_size) { - // Left shift the extra values in the buffer by the batch_size. - const auto buffer_skipped = SkipInBuffer(batch_size); - - // Skip the rest in the decoder - const auto batch_to_skip = batch_size - buffer_skipped; - return Base::Skip(batch_to_skip) + buffer_skipped; -} - -auto StackedLevelDecoder::CountUpTo(int16_t value, - int32_t batch_size) -> CountUpToResult { - // Count in the decoded values - const auto buffer_batch = std::min(batch_size, extra_decoded()); - const auto buffer_count = static_cast( - std::count(extra_start(), extra_start() + buffer_batch, value)); - [[maybe_unused]] const auto buffer_skipped = SkipInBuffer(batch_size); - ARROW_DCHECK_EQ(buffer_skipped, buffer_batch); - - // Count the rest in the decoder. - const auto batch_to_count = batch_size - buffer_batch; - auto out = Base::CountUpTo(value, batch_to_count); - - out.matching_count += buffer_count; - out.processed_count += buffer_batch; - return out; -} - /************************** * ColumnReaderImplBase * **************************/ From 4ee91a6286e0f6fa26dabd6f0026c2b7bfd0befc Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 10:41:21 +0200 Subject: [PATCH 19/33] Remove mark_as_written --- cpp/src/parquet/column_reader.cc | 53 ++++++++++++++++---------------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index b95baaddd5a..da4a4492857 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -784,28 +784,24 @@ class ValueSinkBuffer : private ValueSinkCursor { value_type* data() const { return values_->mutable_data_as(); } - // TODO can we remove that? - value_type* write_start() { return data() + values_count(); } - - // TODO can we remove that? - void mark_values_as_written(int64_t extra_values) { - set_values_count(values_count() + extra_values); - } - void delete_back(int64_t count) { ARROW_DCHECK_LE(count, values_count()); set_values_count(values_count() - count); } [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { - return decoder.Decode(write_start(), batch_size); + const auto decoded = decoder.Decode(write_start(), batch_size); + set_values_count(values_count() + batch_size); + return decoded; } [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, int32_t null_count, const uint8_t* valid_bits, int64_t valid_bits_offset) { - return decoder.DecodeSpaced(write_start(), batch_size, null_count, valid_bits, - valid_bits_offset); + const auto decoded = decoder.DecodeSpaced(write_start(), batch_size, null_count, + valid_bits, valid_bits_offset); + set_values_count(values_count() + batch_size); + return decoded; } std::shared_ptr ReleaseValues(MemoryPool* pool) { @@ -844,6 +840,8 @@ class ValueSinkBuffer : private ValueSinkCursor { } return bytes; } + + value_type* write_start() { return data() + values_count(); } }; /************************** @@ -2256,11 +2254,10 @@ int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { ARROW_DCHECK_GE(values_to_read, 0); ARROW_DCHECK_GE(null_count, 0); + // The values have already been accounted for by ReadValuesDense/Spaced. if (read_dense_for_nullable_) { - value_sink_.mark_values_as_written(values_to_read); ARROW_DCHECK_EQ(null_count, 0); } else { - value_sink_.mark_values_as_written(values_to_read + null_count); null_count_ += null_count; } // Total values, including null spaces, if any @@ -2448,8 +2445,6 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records std::min(clamp_to(num_records - records_read), this->available_values_current_page()); ReadValuesDense(batch_size); - - value_sink_.mark_values_as_written(batch_size); this->ConsumeBufferedValues(batch_size); records_read += batch_size; @@ -2511,8 +2506,6 @@ class ArrayValuesSink : private ValueSinkCursor { value_type* data() const { return nullptr; } - value_type* write_start() { return nullptr; } - void mark_values_as_written(int64_t extra_values) { set_values_count(values_count() + extra_values); } @@ -2533,7 +2526,9 @@ class ArrayValuesSink : private ValueSinkCursor { [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { // TODO once we have a validity sink: we can reset the validity to save some space - return decoder.DecodeArrowNonNull(batch_size, &builder_); + const int64_t decoded = decoder.DecodeArrowNonNull(batch_size, &builder_); + mark_values_as_written(batch_size); + return decoded; } [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, @@ -2542,6 +2537,7 @@ class ArrayValuesSink : private ValueSinkCursor { // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = decoder.DecodeArrow(batch_size, null_count, valid_bits, valid_bits_offset, &builder_); + mark_values_as_written(batch_size); // `DecodeArrow` only counts the non-null values, but the caller expects the // number of values written, nulls included. ARROW_DCHECK_EQ(decoded, batch_size - null_count); @@ -2740,9 +2736,7 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { value_type* data() const { return nullptr; } - void mark_values_as_written(int64_t extra_values) { - set_values_count(values_count() + extra_values); - } + void mark_values_as_written(int64_t extra_values) {} std::shared_ptr ReleaseValues(MemoryPool* /* pool */) { return nullptr; @@ -2754,18 +2748,22 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { [[nodiscard]] int64_t ReadValuesDense(auto& decoder, int32_t batch_size) { // TODO once we have a validity sink: we can reset the validity to save some space - if (auto* dict_decoder = dynamic_cast(&decoder)) { - MaybeWriteNewDictionary(*dict_decoder); - return dict_decoder->DecodeIndices(batch_size, builder_.get()); - } - return decoder.DecodeArrowNonNull(batch_size, builder_.get()); + const int64_t decoded = [&]() -> int64_t { + if (auto* dict_decoder = dynamic_cast(&decoder)) { + MaybeWriteNewDictionary(*dict_decoder); + return dict_decoder->DecodeIndices(batch_size, builder_.get()); + } + return decoder.DecodeArrowNonNull(batch_size, builder_.get()); + }(); + set_values_count(values_count() + batch_size); + return decoded; } [[nodiscard]] int64_t ReadValuesSpaced(auto& decoder, int32_t batch_size, int32_t null_count, const uint8_t* valid_bits, int64_t valid_bits_offset) { // TODO once we have a validity sink: we can reset the validity to save some space - const int64_t decoded = [&] { + const int64_t decoded = [&]() -> int64_t { if (auto* dict_decoder = dynamic_cast(&decoder)) { MaybeWriteNewDictionary(*dict_decoder); return dict_decoder->DecodeIndicesSpaced(batch_size, null_count, valid_bits, @@ -2774,6 +2772,7 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { return decoder.DecodeArrow(batch_size, null_count, valid_bits, valid_bits_offset, builder_.get()); }(); + set_values_count(values_count() + batch_size); // The decoders only count the non-null values, but the caller expects the // number of values written, nulls included. ARROW_DCHECK_EQ(decoded, batch_size - null_count); From b8e89dd94e7e21fd87a4356a37da0dda3aedc318 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 10:42:00 +0200 Subject: [PATCH 20/33] Make some of ColumnReaderImplBase private --- cpp/src/parquet/column_reader.cc | 80 +++++++++++++------------------- 1 file changed, 33 insertions(+), 47 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index da4a4492857..902a6982e47 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -922,17 +922,19 @@ class ColumnReaderImplBase { ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, LvlDec def_levels_decoder, LvlDec rep_levels_decoder) - : def_levels_decoder_(std::move(def_levels_decoder)), - rep_levels_decoder_(std::move(rep_levels_decoder)), - descr_(descr), + : descr_(descr), pool_(pool), - current_decoder_(pool) {} + current_decoder_(pool), + def_levels_decoder_(std::move(def_levels_decoder)), + rep_levels_decoder_(std::move(rep_levels_decoder)) {} virtual ~ColumnReaderImplBase() = default; - protected: + private: + // Declared here because `decoders_` below refers to it. using DecoderType = TypedDecoder; + protected: int32_t ReadDefinitionLevels(int32_t batch_size, int16_t* levels) { if (max_def_level() == 0) { return 0; @@ -947,21 +949,8 @@ class ColumnReaderImplBase { return rep_levels_decoder_.Decode(batch_size, levels); } - LvlDec def_levels_decoder_; - LvlDec rep_levels_decoder_; const ColumnDescriptor* descr_; std::unique_ptr pager_; - std::shared_ptr current_page_; - // The total number of values stored in the data page. This is the maximum of - // the number of encoded definition levels or encoded values. For - // non-repeated, required columns, this is equal to the number of encoded - // 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_ = 0; - // The number of values from the current data page that have been decoded - // into memory or skipped over. - int64_t num_decoded_values_ = 0; ::arrow::MemoryPool* pool_; SkippableTypedDecoder current_decoder_; Encoding::type current_encoding_ = Encoding::UNKNOWN; @@ -981,25 +970,8 @@ class ColumnReaderImplBase { // @returns: the number of values read into the out buffer int64_t ReadValues(int64_t batch_size, T* out); - // Read up to batch_size values from the current data page into the - // pre-allocated memory T*, leaving spaces for null entries according - // to the def_levels. - // - // @returns: the number of values read into the out buffer - int64_t ReadValuesSpaced(int64_t batch_size, T* out, int64_t null_count, - uint8_t* valid_bits, int64_t valid_bits_offset); - bool HasNextInternal(); - // Advance to the next data page - bool ReadNewPage(); - - void ConfigureDictionary(const DictionaryPage* page); - - // Get a decoder object for this page or create a new decoder if this is the - // first page with this encoding. - void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size); - // Available values in the current data page, value includes repeated values // and nulls. int32_t available_values_current_page() const { @@ -1016,13 +988,37 @@ class ColumnReaderImplBase { void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } + int64_t Skip(int64_t num_values_to_skip); + + private: + LvlDec def_levels_decoder_; + LvlDec rep_levels_decoder_; + std::shared_ptr current_page_; + // The total number of values stored in the data page. This is the maximum of + // the number of encoded definition levels or encoded values. For + // non-repeated, required columns, this is equal to the number of encoded + // 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_ = 0; + // The number of values from the current data page that have been decoded + // into memory or skipped over. + int64_t num_decoded_values_ = 0; + + // Advance to the next data page + bool ReadNewPage(); + + void ConfigureDictionary(const DictionaryPage* page); + + // Get a decoder object for this page or create a new decoder if this is the + // first page with this encoding. + void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size); + /// Advance both level decoders by num_levels, without materializing them. /// /// @return The number of values present (non-null) among them, which is the /// number of values to skip in the data decoder. int32_t AdvanceLevels(int32_t num_levels); - - int64_t Skip(int64_t num_values_to_skip); }; /***************************************** @@ -1035,16 +1031,6 @@ int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* o return num_decoded; } -template -int64_t ColumnReaderImplBase::ReadValuesSpaced(int64_t batch_size, T* out, - int64_t null_count, - uint8_t* valid_bits, - int64_t valid_bits_offset) { - return current_decoder_->DecodeSpaced(out, static_cast(batch_size), - static_cast(null_count), valid_bits, - valid_bits_offset); -} - template bool ColumnReaderImplBase::HasNextInternal() { // Either there is no data page available yet, or the data page has been From 14d7007e4c1dd4427bfa8a5652fbf47a07614677 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 11:37:15 +0200 Subject: [PATCH 21/33] Fix horrible dictionnary handling --- cpp/src/parquet/column_reader.cc | 55 +++++++++++++++----------------- 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 902a6982e47..26ca1aafcc4 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -784,6 +784,8 @@ class ValueSinkBuffer : private ValueSinkCursor { value_type* data() const { return values_->mutable_data_as(); } + void OnNewDictionary(auto& /* decoder */) {} + void delete_back(int64_t count) { ARROW_DCHECK_LE(count, values_count()); set_values_count(values_count() - count); @@ -954,9 +956,6 @@ class ColumnReaderImplBase { ::arrow::MemoryPool* pool_; SkippableTypedDecoder current_decoder_; Encoding::type current_encoding_ = Encoding::UNKNOWN; - /// Flag to signal when a new dictionary has been set, for the benefit of - /// DictionaryRecordReader - bool new_dictionary_ = false; // The exposed encoding ExposedEncoding exposed_encoding_ = ExposedEncoding::NO_ENCODING; // Map of encoding type to the respective decoder object. For example, a @@ -972,6 +971,12 @@ class ColumnReaderImplBase { bool HasNextInternal(); + /// Called when a dictionary page has been read and its decoder set up. + /// + /// Readers accumulating dictionary-encoded values need to know about it, as + /// the indices they have read so far refer to the previous dictionary. + virtual void OnNewDictionary(DictDecoder& /* decoder */) {} + // Available values in the current data page, value includes repeated values // and nulls. int32_t available_values_current_page() const { @@ -1109,13 +1114,13 @@ void ColumnReaderImplBase::ConfigureDictionary( std::unique_ptr> decoder = MakeDictDecoder(descr_, pool_); decoder->SetDict(dictionary.get()); + OnNewDictionary(*decoder); decoders_[encoding] = std::unique_ptr(dynamic_cast(decoder.release())); } else { ParquetException::NYI("only plain dictionary encoding has been implemented"); } - new_dictionary_ = true; current_decoder_.SetDecoder(decoders_[encoding].get()); ARROW_DCHECK(current_decoder_); } @@ -1592,6 +1597,10 @@ class TypedRecordReader : public ColumnReaderImplBase, // Dictionary decoders must be reset when advancing row groups void ResetDecoders() { this->decoders_.clear(); } + void OnNewDictionary(DictDecoder& decoder) final { + value_sink_.OnNewDictionary(decoder); + } + void ReadValuesSpaced(int32_t values_with_nulls, int32_t null_count); void ReadValuesDense(int32_t values_to_read); @@ -2371,6 +2380,10 @@ class RequiredTypedRecordReader : public ColumnReaderImplBasedecoders_.clear(); } + void OnNewDictionary(DictDecoder& decoder) final { + value_sink_.OnNewDictionary(decoder); + } + void DebugPrintState() final; protected: @@ -2510,6 +2523,8 @@ class ArrayValuesSink : private ValueSinkCursor { void ResetValues() { ValueSinkCursor::operator=({}); } + void OnNewDictionary(auto& /* decoder */) {} + [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = decoder.DecodeArrowNonNull(batch_size, &builder_); @@ -2736,7 +2751,6 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = [&]() -> int64_t { if (auto* dict_decoder = dynamic_cast(&decoder)) { - MaybeWriteNewDictionary(*dict_decoder); return dict_decoder->DecodeIndices(batch_size, builder_.get()); } return decoder.DecodeArrowNonNull(batch_size, builder_.get()); @@ -2751,7 +2765,6 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = [&]() -> int64_t { if (auto* dict_decoder = dynamic_cast(&decoder)) { - MaybeWriteNewDictionary(*dict_decoder); return dict_decoder->DecodeIndicesSpaced(batch_size, null_count, valid_bits, valid_bits_offset, builder_.get()); } @@ -2765,11 +2778,12 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { return decoded + null_count; } - /// Bind the reader's flag signalling that a new dictionary page was read. - /// - /// The flag belongs to the reader base class, which is only constructed - /// once this sink has been handed over to it, hence the late binding. - void BindNewDictionaryFlag(bool* new_dictionary) { new_dictionary_ = new_dictionary; } + void OnNewDictionary(auto& decoder) { + // The indices accumulated so far refer to the previous dictionary + FlushBuilder(); + builder_->ResetFull(); + decoder.InsertDictionary(builder_.get()); + } /// Finish the builder and return all the chunks accumulated so far. std::shared_ptr<::arrow::ChunkedArray> FlushChunks() { @@ -2783,7 +2797,6 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { std::unique_ptr builder_; std::vector> result_chunks_; - bool* new_dictionary_ = nullptr; void FlushBuilder() { if (builder_->length() > 0) { @@ -2795,18 +2808,6 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { builder_->Reset(); } } - - void MaybeWriteNewDictionary(BinaryDictDecoder& decoder) { - ARROW_DCHECK_NE(new_dictionary_, nullptr); - if (*new_dictionary_) { - /// If there is a new dictionary, we may need to flush the builder, then - /// insert the new dictionary values - FlushBuilder(); - builder_->ResetFull(); - decoder.InsertDictionary(builder_.get()); - *new_dictionary_ = false; - } - } }; template @@ -2841,7 +2842,6 @@ class ByteArrayDictionaryRecordReader final requires(!kRequired) : Base(descr, leaf_info, pool, read_dense_for_nullable, ValueSink(pool)) { ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); - BindNewDictionaryFlag(); } ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, @@ -2851,7 +2851,6 @@ class ByteArrayDictionaryRecordReader final ARROW_DCHECK_EQ(descr->physical_type(), Type::BYTE_ARRAY); ARROW_DCHECK_EQ(descr->max_definition_level(), 0); ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); - BindNewDictionaryFlag(); } std::shared_ptr<::arrow::ChunkedArray> GetResult() override { @@ -2860,10 +2859,6 @@ class ByteArrayDictionaryRecordReader final private: using ValueSink = typename byte_array_dictionary_record_reader::ValueSink; - - void BindNewDictionaryFlag() { - this->value_sink().BindNewDictionaryFlag(&this->new_dictionary_); - } }; std::shared_ptr MakeByteArrayRecordReader( From 4858c42d630304fe69ec0c502270a18f478bc22a Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Mon, 20 Jul 2026 18:22:26 +0200 Subject: [PATCH 22/33] Simplify and rename ColumnChunkReader --- cpp/src/parquet/column_reader.cc | 334 +++++++++++++++++-------------- 1 file changed, 186 insertions(+), 148 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 26ca1aafcc4..2f00a9f6637 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -81,15 +81,22 @@ constexpr int64_t kSkipScratchBatchSize = 1024; // Throws exception if number_decoded does not match expected. inline void CheckNumberDecoded(int64_t number_decoded, int64_t expected) { if (ARROW_PREDICT_FALSE(number_decoded != expected)) { - ParquetException::EofException("Decoded values " + std::to_string(number_decoded) + - " does not match expected " + - std::to_string(expected)); + auto msg = std::format("Decoded values {} does not match expected {}", number_decoded, + expected); + ParquetException::EofException(msg); } } constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = "Number of decoded rep / def levels do not match num_values in page header"; +template +constexpr T clamp_to(U val) { + constexpr U kMax = std::numeric_limits::max(); + constexpr U kMin = std::numeric_limits::min(); + return static_cast(std::clamp(val, kMin, kMax)); +} + } // namespace /****************** @@ -846,9 +853,9 @@ class ValueSinkBuffer : private ValueSinkCursor { value_type* write_start() { return data() + values_count(); } }; -/************************** - * ColumnReaderImplBase * - **************************/ +/*********************** + * ColumnChunkReader * + ***********************/ /// Initialize repetition and definition level decoders on the given data page. /// @@ -916,26 +923,50 @@ int64_t InitializeV2Levels(const DataPageV2& page, LvlDec& def_dec, LvlDec& rep_ return total_levels_length; } -/// Impl base class for TypedColumnReader and RecordReader -template -class ColumnReaderImplBase { +/// Traits of a concrete column chunk reader, for use by `ColumnChunkReader`. +template +struct reader_trait; + +/// Read through the multiple pages of a column chunk. +template +class ColumnChunkReader { public: - using T = typename DType::c_type; + using DType = typename reader_trait::DType; + using LvlDec = typename reader_trait::level_decoder; + using value_type = typename DType::c_type; - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, - LvlDec def_levels_decoder, LvlDec rep_levels_decoder) + ColumnChunkReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, + LvlDec def_levels_decoder, LvlDec rep_levels_decoder) : descr_(descr), pool_(pool), current_decoder_(pool), def_levels_decoder_(std::move(def_levels_decoder)), rep_levels_decoder_(std::move(rep_levels_decoder)) {} - virtual ~ColumnReaderImplBase() = default; + void SetPageReader(std::unique_ptr reader) { + pager_ = std::move(reader); + // Dictionary decoders must be reset when advancing row groups + decoders_.clear(); + } + + bool HasPageReader() const { return pager_ != nullptr; } + + /// Return true if there is more data. + /// + /// If the current page is exhausted, it will process more pages until some data + /// page is found. + bool ProcessToMoreData(); + + /// Check the encoding of the current page or throw an exception. + void CheckEncodingIs(Encoding::type encoding); private: // Declared here because `decoders_` below refers to it. using DecoderType = TypedDecoder; + Derived& derived() { return static_cast(*this); } + const Derived& derived() const { return static_cast(*this); } + protected: int32_t ReadDefinitionLevels(int32_t batch_size, int16_t* levels) { if (max_def_level() == 0) { @@ -952,39 +983,14 @@ class ColumnReaderImplBase { } const ColumnDescriptor* descr_; - std::unique_ptr pager_; ::arrow::MemoryPool* pool_; SkippableTypedDecoder current_decoder_; - Encoding::type current_encoding_ = Encoding::UNKNOWN; - // The exposed encoding - ExposedEncoding exposed_encoding_ = ExposedEncoding::NO_ENCODING; - // Map of encoding type to the respective decoder object. For example, a - // column chunk's data pages may include both dictionary-encoded and - // plain-encoded data. - std::unordered_map> decoders_; - // Read up to batch_size values from the current data page into the - // pre-allocated memory T* - // - // @returns: the number of values read into the out buffer - int64_t ReadValues(int64_t batch_size, T* out); - - bool HasNextInternal(); - - /// Called when a dictionary page has been read and its decoder set up. - /// - /// Readers accumulating dictionary-encoded values need to know about it, as - /// the indices they have read so far refer to the previous dictionary. - virtual void OnNewDictionary(DictDecoder& /* decoder */) {} - - // Available values in the current data page, value includes repeated values - // and nulls. + // Available values in the current data page, value includes repeated values and nulls. int32_t available_values_current_page() const { - // The number of values in a page fits inside an int32_t according to the spec. - const int64_t out = num_buffered_values_ - num_decoded_values_; + const int32_t out = num_buffered_values_ - num_decoded_values_; ARROW_DCHECK_GE(out, 0); - ARROW_DCHECK_LE(out, std::numeric_limits::max()); - return static_cast(out); + return out; } int16_t max_def_level() const; @@ -998,6 +1004,11 @@ class ColumnReaderImplBase { private: LvlDec def_levels_decoder_; LvlDec rep_levels_decoder_; + // Map of encoding type to the respective decoder object. For example, a + // column chunk's data pages may include both dictionary-encoded and + // plain-encoded data. + std::unordered_map> decoders_; + std::unique_ptr pager_; std::shared_ptr current_page_; // The total number of values stored in the data page. This is the maximum of // the number of encoded definition levels or encoded values. For @@ -1005,10 +1016,11 @@ 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_ = 0; + int32_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_ = 0; + int32_t num_decoded_values_ = 0; + Encoding::type current_encoding_ = Encoding::UNKNOWN; // Advance to the next data page bool ReadNewPage(); @@ -1026,18 +1038,22 @@ class ColumnReaderImplBase { int32_t AdvanceLevels(int32_t num_levels); }; -/***************************************** - * ColumnReaderImplBase Implementation * - *****************************************/ +/************************************** + * ColumnChunkReader Implementation * + **************************************/ -template -int64_t ColumnReaderImplBase::ReadValues(int64_t batch_size, T* out) { - int64_t num_decoded = current_decoder_->Decode(out, static_cast(batch_size)); - return num_decoded; +template +void ColumnChunkReader::CheckEncodingIs(Encoding::type encoding) { + if (current_encoding_ != encoding) { + auto msg = + std::format("Unexpected data page encoding. Expected {}, got {}", + EncodingToString(encoding), EncodingToString(current_encoding_)); + throw ParquetException(msg); + } } -template -bool ColumnReaderImplBase::HasNextInternal() { +template +bool ColumnChunkReader::ProcessToMoreData() { // Either there is no data page available yet, or the data page has been // exhausted if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { @@ -1048,8 +1064,8 @@ bool ColumnReaderImplBase::HasNextInternal() { return true; } -template -bool ColumnReaderImplBase::ReadNewPage() { +template +bool ColumnChunkReader::ReadNewPage() { // Loop until we find the next data page. while (true) { current_page_ = pager_->NextPage(); @@ -1087,9 +1103,8 @@ bool ColumnReaderImplBase::ReadNewPage() { return true; } -template -void ColumnReaderImplBase::ConfigureDictionary( - const DictionaryPage* page) { +template +void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) { int encoding = static_cast(page->encoding()); if (page->encoding() == Encoding::PLAIN_DICTIONARY || page->encoding() == Encoding::PLAIN) { @@ -1114,7 +1129,7 @@ void ColumnReaderImplBase::ConfigureDictionary( std::unique_ptr> decoder = MakeDictDecoder(descr_, pool_); decoder->SetDict(dictionary.get()); - OnNewDictionary(*decoder); + derived().OnNewDictionary(*decoder); decoders_[encoding] = std::unique_ptr(dynamic_cast(decoder.release())); } else { @@ -1125,9 +1140,9 @@ void ColumnReaderImplBase::ConfigureDictionary( ARROW_DCHECK(current_decoder_); } -template -void ColumnReaderImplBase::InitializeDataDecoder( - const DataPage& page, int64_t levels_byte_size) { +template +void ColumnChunkReader::InitializeDataDecoder(const DataPage& page, + int64_t levels_byte_size) { const uint8_t* buffer = page.data() + levels_byte_size; const int64_t data_size = page.size() - levels_byte_size; @@ -1172,18 +1187,18 @@ void ColumnReaderImplBase::InitializeDataDecoder( static_cast(data_size)); } -template -int16_t ColumnReaderImplBase::max_def_level() const { +template +int16_t ColumnChunkReader::max_def_level() const { return def_levels_decoder_.max_level(); } -template -int16_t ColumnReaderImplBase::max_rep_level() const { +template +int16_t ColumnChunkReader::max_rep_level() const { return rep_levels_decoder_.max_level(); } -template -int32_t ColumnReaderImplBase::AdvanceLevels(int32_t num_levels) { +template +int32_t ColumnChunkReader::AdvanceLevels(int32_t num_levels) { int max_count = num_levels; // Advance the definition levels, counting how many correspond to present // (non-null) values that must be skipped in the data decoder. @@ -1199,11 +1214,11 @@ int32_t ColumnReaderImplBase::AdvanceLevels(int32_t num_levels) { return max_count; } -template -int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { +template +int64_t ColumnChunkReader::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 && HasNextInternal()) { + while (values_to_skip > 0 && ProcessToMoreData()) { // If the number of values to skip is more than the number of undecoded values, skip // the whole Page without decoding levels or values. const int64_t available_values = this->available_values_current_page(); @@ -1226,44 +1241,50 @@ int64_t ColumnReaderImplBase::Skip(int64_t num_values_to_skip) { return num_values_to_skip - values_to_skip; } -template -constexpr T clamp_to(U val) { - constexpr U kMax = std::numeric_limits::max(); - constexpr U kMin = std::numeric_limits::min(); - return static_cast(std::clamp(val, kMin, kMax)); -} +/*************************** + * TypedColumnReaderImpl * + ***************************/ -// ---------------------------------------------------------------------- -// TypedColumnReader implementations +template +class TypedColumnReaderImpl; + +template +struct reader_trait> { + using DType = D; + using level_decoder = LevelDecoder; +}; template class TypedColumnReaderImpl : public TypedColumnReader, - public ColumnReaderImplBase { + public ColumnChunkReader> { public: using T = typename DType::c_type; TypedColumnReaderImpl(const ColumnDescriptor* descr, std::unique_ptr pager, ::arrow::MemoryPool* pool) - : ColumnReaderImplBase( + : ColumnChunkReader>( descr, pool, LevelDecoder(descr->max_definition_level()), LevelDecoder(descr->max_repetition_level())) { - this->pager_ = std::move(pager); + this->SetPageReader(std::move(pager)); } - bool HasNext() override { return this->HasNextInternal(); } + bool HasNext() override { return this->ProcessToMoreData(); } + + /// Called by ColumnChunkReader. + void OnNewDictionary(DictDecoder& /* decoder */) {} int64_t ReadBatch(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, T* values, int64_t* values_read) override; int64_t Skip(int64_t num_values_to_skip) override { - return ColumnReaderImplBase::Skip(num_values_to_skip); + return ColumnChunkReader>::Skip(num_values_to_skip); } Type::type type() const override { return this->descr_->physical_type(); } const ColumnDescriptor* descr() const override { return this->descr_; } - ExposedEncoding GetExposedEncoding() override { return this->exposed_encoding_; }; + ExposedEncoding GetExposedEncoding() override { return exposed_encoding_; }; int64_t ReadBatchWithDictionary(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, int32_t* indices, @@ -1272,10 +1293,13 @@ class TypedColumnReaderImpl : public TypedColumnReader, protected: void SetExposedEncoding(ExposedEncoding encoding) override { - this->exposed_encoding_ = encoding; + exposed_encoding_ = encoding; } private: + // The exposed encoding + ExposedEncoding exposed_encoding_ = ExposedEncoding::NO_ENCODING; + // 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) { @@ -1297,8 +1321,8 @@ class TypedColumnReaderImpl : public TypedColumnReader, // ReadLevelsInCurrentPage will throw exception when any num-levels read is not // equal to the number of the levels can be read. void ReadLevelsInCurrentPage(int32_t batch_size, int16_t* def_levels, - int16_t* rep_levels, int64_t* num_def_levels, - int64_t* non_null_values_to_read) { + int16_t* rep_levels, int32_t* num_def_levels, + int32_t* non_null_values_to_read) { 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) { @@ -1347,12 +1371,7 @@ int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( } // Verify the current data page is dictionary encoded. - if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { - std::stringstream ss; - ss << "Data page is not dictionary encoded. Encoding: " - << EncodingToString(this->current_encoding_); - throw ParquetException(ss.str()); - } + this->CheckEncodingIs(Encoding::RLE_DICTIONARY); // Get dictionary pointer and length. if (has_dict_output) { @@ -1360,8 +1379,8 @@ int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( } // Similar logic as ReadValues to get def levels and rep levels. - int64_t num_def_levels = 0; - int64_t indices_to_read = 0; + int32_t num_def_levels = 0; + int32_t indices_to_read = 0; ReadLevelsInCurrentPage(batch_size, def_levels, rep_levels, &num_def_levels, &indices_to_read); @@ -1396,9 +1415,9 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size_64, // TODO(wesm): keep reading data pages until batch_size is reached, or the // row group is finished - int64_t num_def_levels = 0; + int32_t num_def_levels = 0; // Number of non-null values to read within `num_def_levels`. - int64_t non_null_values_to_read = 0; + int32_t non_null_values_to_read = 0; ReadLevelsInCurrentPage(batch_size, def_levels, rep_levels, &num_def_levels, &non_null_values_to_read); // Should not return more values than available in the current data page, @@ -1408,7 +1427,7 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size_64, throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } if (non_null_values_to_read != 0) { - *values_read = this->ReadValues(non_null_values_to_read, values); + *values_read = this->current_decoder_->Decode(values, non_null_values_to_read); } else { *values_read = 0; } @@ -1482,11 +1501,30 @@ concept can_cout = requires(std::ostream& os, const T& value) { ***********************/ template -class TypedRecordReader : public ColumnReaderImplBase, - virtual public RecordReader { +class TypedRecordReader; + +// `reader_trait` can only be specialized in the namespace it is declared in. +} // namespace +} // namespace internal + +namespace { +template +struct reader_trait> { + using DType = D; + using level_decoder = LevelDecoder; +}; +} // namespace + +namespace internal { +namespace { + +template +class TypedRecordReader + : public ColumnChunkReader>, + virtual public RecordReader { public: using T = typename DType::c_type; - using Base = ColumnReaderImplBase; + using Base = ColumnChunkReader>; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable, ValueSink value_sink) @@ -1590,14 +1628,13 @@ class TypedRecordReader : public ColumnReaderImplBase, void SetPageReader(std::unique_ptr reader) override; - bool HasMoreData() const override { return this->pager_ != nullptr; } + bool HasMoreData() const override { + return Base::HasPageReader(); // Surprising legacy behaviour + } const ColumnDescriptor* descr() const override { return this->descr_; } - // Dictionary decoders must be reset when advancing row groups - void ResetDecoders() { this->decoders_.clear(); } - - void OnNewDictionary(DictDecoder& decoder) final { + void OnNewDictionary(DictDecoder& decoder) { value_sink_.OnNewDictionary(decoder); } @@ -1689,19 +1726,12 @@ class TypedRecordReader : public ColumnReaderImplBase, template const void* TypedRecordReader::ReadDictionary(int32_t* dictionary_length) { - if (!this->current_decoder_ && !this->HasNextInternal()) { + if (!this->current_decoder_ && !this->ProcessToMoreData()) { *dictionary_length = 0; return nullptr; } - // Verify the current data page is dictionary encoded. The current_encoding_ should - // have been set as RLE_DICTIONARY if the page encoding is RLE_DICTIONARY or - // PLAIN_DICTIONARY. - if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { - std::stringstream ss; - ss << "Data page is not dictionary encoded. Encoding: " - << EncodingToString(this->current_encoding_); - throw ParquetException(ss.str()); - } + // Verify the current data page is dictionary encoded. + this->CheckEncodingIs(Encoding::RLE_DICTIONARY); auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); @@ -1725,7 +1755,7 @@ int64_t TypedRecordReader::ReadRecords(int64_t num_records) { // enough records while (!at_record_start_ || records_read < num_records) { // Is there more data to read in this row group? - if (!this->HasNextInternal()) { + if (!this->ProcessToMoreData()) { if (!at_record_start_) { // We ended the row group while inside a record that we haven't seen // the end of yet. So increment the record count for the last record in @@ -1884,7 +1914,7 @@ int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_records while (!at_record_start_ || skipped_records < num_records) { // Is there more data to read in this row group? // HasNextInternal() will advance to the next page if necessary. - if (!this->HasNextInternal()) { + if (!this->ProcessToMoreData()) { if (!at_record_start_) { // We ended the row group while inside a record that we haven't seen // the end of yet. So increment the record count for the last record @@ -2106,8 +2136,7 @@ void TypedRecordReader::Reset() { template void TypedRecordReader::SetPageReader(std::unique_ptr reader) { at_record_start_ = true; - this->pager_ = std::move(reader); - ResetDecoders(); + Base::SetPageReader(std::move(reader)); } template @@ -2317,13 +2346,34 @@ void TypedRecordReader::ResetValues() { *******************************/ // TODO can we reduce some code share with TypedRecordREader ? +template +class RequiredTypedRecordReader; + +// `reader_trait` can only be specialized in the namespace it is declared in. +} // namespace +} // namespace internal + +namespace { +template +struct reader_trait> { + using DType = D; + using level_decoder = LevelDecoder; +}; +} // namespace + +namespace internal { +namespace { + template , bool kReadDictionary = false> -class RequiredTypedRecordReader : public ColumnReaderImplBase, - virtual public RecordReader { +class RequiredTypedRecordReader + : public ColumnChunkReader< + RequiredTypedRecordReader>, + virtual public RecordReader { public: using T = typename DType::c_type; - using Base = ColumnReaderImplBase; + using Base = + ColumnChunkReader>; RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool, ValueSink value_sink) @@ -2371,16 +2421,17 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase reader) final; + void SetPageReader(std::unique_ptr reader) final { + return Base::SetPageReader(std::move(reader)); + } - bool HasMoreData() const final { return this->pager_ != nullptr; } + bool HasMoreData() const override { + return Base::HasPageReader(); // Surprising legacy behaviour + } const ColumnDescriptor* descr() const final { return this->descr_; } - // Dictionary decoders must be reset when advancing row groups - void ResetDecoders() { this->decoders_.clear(); } - - void OnNewDictionary(DictDecoder& decoder) final { + void OnNewDictionary(DictDecoder& decoder) { value_sink_.OnNewDictionary(decoder); } @@ -2405,19 +2456,12 @@ class RequiredTypedRecordReader : public ColumnReaderImplBase const void* RequiredTypedRecordReader::ReadDictionary( int32_t* dictionary_length) { - if (!this->current_decoder_ && !this->HasNextInternal()) { + if (!this->current_decoder_ && !this->ProcessToMoreData()) { *dictionary_length = 0; return nullptr; } - // Verify the current data page is dictionary encoded. The current_encoding_ should - // have been set as RLE_DICTIONARY if the page encoding is RLE_DICTIONARY or - // PLAIN_DICTIONARY. - if (this->current_encoding_ != Encoding::RLE_DICTIONARY) { - std::stringstream ss; - ss << "Data page is not dictionary encoded. Encoding: " - << EncodingToString(this->current_encoding_); - throw ParquetException(ss.str()); - } + // Verify the current data page is dictionary encoded. + this->CheckEncodingIs(Encoding::RLE_DICTIONARY); auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); @@ -2436,7 +2480,7 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records do { // Is there more data to read in this row group? - if (!this->HasNextInternal()) { + if (!this->ProcessToMoreData()) { break; } @@ -2466,13 +2510,6 @@ void RequiredTypedRecordReader::Reset() { } } -template -void RequiredTypedRecordReader::SetPageReader( - std::unique_ptr reader) { - this->pager_ = std::move(reader); - ResetDecoders(); -} - template void RequiredTypedRecordReader::DebugPrintState() { std::cout << "values: "; @@ -2804,7 +2841,8 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { PARQUET_THROW_NOT_OK(builder_->Finish(&chunk)); result_chunks_.emplace_back(std::move(chunk)); - // Also clears the dictionary memo table + // Partial reset: the dictionary memo table is kept, so that the indices + // appended to the next chunk keep referring to the same values. builder_->Reset(); } } From 28c60b3c54e446dfd4ca72dbb2c063dd66b5997f Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 21 Jul 2026 15:18:28 +0200 Subject: [PATCH 23/33] Add RleBitPackedToBitmapDecoder::CountUpTo --- cpp/src/arrow/util/rle_bitmap_internal.h | 101 +++++++++++++++---- cpp/src/arrow/util/rle_bitmap_test.cc | 119 ++++++++++++++++++++++- 2 files changed, 195 insertions(+), 25 deletions(-) diff --git a/cpp/src/arrow/util/rle_bitmap_internal.h b/cpp/src/arrow/util/rle_bitmap_internal.h index 0667a9f0a8e..f1395a008a9 100644 --- a/cpp/src/arrow/util/rle_bitmap_internal.h +++ b/cpp/src/arrow/util/rle_bitmap_internal.h @@ -112,6 +112,18 @@ class RleRunToBitmapDecoder { return n_vals; } + /// Advance and count the number of occurrences of `value`. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of elements that were processed. + RleCountUpToResult CountUpTo(bool value, rle_size_t batch_size) { + const auto n_vals = Advance(batch_size); + return { + .matching_count = n_vals * (value == this->value()), + .processed_count = n_vals, + }; + } + /// Read the next value into `out` and return false if there are no more. [[nodiscard]] bool Get(BitmapSpanMut out) { return GetBatch(out, 1) == 1; } @@ -229,6 +241,19 @@ class BitPackedRunToBitmapDecoder { return n_vals; } + /// Advance and count the number of occurrences of `value`. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of elements that were processed. + RleCountUpToResult CountUpTo(bool value, rle_size_t batch_size) { + const auto n_vals = GetAdvanceCapacity(batch_size); + const auto set_count = static_cast(arrow::internal::CountSetBits( + unread_values_ptr(), unread_values_bit_offset(), n_vals)); + const auto matching_count = value ? set_count : n_vals - set_count; + ARROW_UNUSED(Advance(n_vals)); + return {.matching_count = matching_count, .processed_count = n_vals}; + } + /// Get the next value and return false if there are no more. [[nodiscard]] bool Get(BitmapSpanMut out) { return GetBatch(out, 1) == 1; } @@ -287,6 +312,14 @@ class RleBitPackedToBitmapDecoder { /// values left or if an error occurred. [[nodiscard]] rle_size_t GetBatch(BitmapSpanMut out, rle_size_t batch_size); + /// Advance and count the number of occurrences of `value`. + /// + /// The count is limited to at most the next `batch_size` items. Fewer than + /// `batch_size` elements may be processed if there are not enough values left + /// or if an error occurred. + /// @return The matching value count and number of elements that were processed. + RleCountUpToResult CountUpTo(bool value, rle_size_t batch_size); + private: RleBitPackedParser parser_ = {}; std::variant decoder_ = {}; @@ -296,10 +329,15 @@ 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 a + /// 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); }; /************************************************ @@ -320,24 +358,26 @@ 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); } @@ -345,17 +385,14 @@ inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, parser_.ParseWithCallable([&](auto run) { using RunDecoder = RleBitPackedToBitmapDecoderGetDecoder::type; - ARROW_DCHECK_LT(values_read, batch_size); + ARROW_DCHECK_LT(values_processed, batch_size); 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; + 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; } @@ -363,7 +400,31 @@ inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, 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) { + const auto read = decoder.GetBatch(out, run_batch_size); + out = out.NewStartingAt(read); + return read; + }, + batch_size); +} + +inline auto RleBitPackedToBitmapDecoder::CountUpTo(bool value, rle_size_t batch_size) + -> RleCountUpToResult { + rle_size_t matching_count = 0; + const rle_size_t processed_count = ProcessValues( + [value, &matching_count](auto& decoder, rle_size_t run_batch_size) { + const auto result = decoder.CountUpTo(value, run_batch_size); + matching_count += result.matching_count; + return result.processed_count; + }, + batch_size); + return {.matching_count = matching_count, .processed_count = processed_count}; } } // namespace arrow::util diff --git a/cpp/src/arrow/util/rle_bitmap_test.cc b/cpp/src/arrow/util/rle_bitmap_test.cc index 7dc572eb5ad..a2040ea9ad6 100644 --- a/cpp/src/arrow/util/rle_bitmap_test.cc +++ b/cpp/src/arrow/util/rle_bitmap_test.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include #include #include #include @@ -285,6 +286,51 @@ void CheckBitmapDecoder(const typename Decoder::RunType& run, } } +/// Reference count of `value` in the slice `expected[start, start + count)`. +rle_size_t CountValueSlice(const std::vector& expected, rle_size_t start, + rle_size_t count, bool value) { + return static_cast( + std::count(expected.begin() + start, expected.begin() + start + count, value)); +} + +/// Count both boolean values in the run with CountUpTo, `chunk_size` values at a time. +/// +/// Verifies the per-chunk and total matching counts, the processed counts, and that +/// counting past the end yields nothing. +template +void CheckDecoderCountUpTo(const typename Decoder::RunType& run, + const std::vector& expected, rle_size_t chunk_size) { + ARROW_SCOPED_TRACE("chunk_size = ", chunk_size); + const auto n_vals = static_cast(expected.size()); + + for (bool value : {true, false}) { + ARROW_SCOPED_TRACE("value = ", value); + Decoder decoder(run); + + rle_size_t processed = 0; + rle_size_t matching = 0; + while (processed < n_vals) { + const auto want = std::min(chunk_size, n_vals - processed); + const auto result = decoder.CountUpTo(value, want); + EXPECT_EQ(result.processed_count, want) << "at pos " << processed; + ASSERT_GT(result.processed_count, 0) << "at pos " << processed; // break on failure + EXPECT_EQ(result.matching_count, + CountValueSlice(expected, processed, result.processed_count, value)) + << "at pos " << processed; + processed += result.processed_count; + matching += result.matching_count; + EXPECT_EQ(decoder.remaining(), n_vals - processed); + } + EXPECT_EQ(processed, n_vals); + EXPECT_EQ(matching, CountValueSlice(expected, 0, n_vals, value)); + + // Counting past the end processes nothing. + const auto past = decoder.CountUpTo(value, std::max(chunk_size, rle_size_t{1})); + EXPECT_EQ(past.processed_count, 0); + EXPECT_EQ(past.matching_count, 0); + } +} + } // namespace /*************************** @@ -310,6 +356,10 @@ TEST_P(RleRunToBitmapDecoderTest, Decode) { const std::vector expected(count, value); CheckBitmapDecoder(run, expected); + + for (const rle_size_t chunk_size : {1, 7, 8, 9, count, count + 1}) { + CheckDecoderCountUpTo(run, expected, chunk_size); + } } } @@ -345,6 +395,10 @@ TEST_P(BitPackedRunToBitmapDecoderTest, Decode) { const std::vector expected = BitsFromBytes(param.bytes, param.count); CheckBitmapDecoder(run, expected); + + for (const rle_size_t chunk_size : {1, 3, 7, 8, 9, param.count, param.count + 1}) { + CheckDecoderCountUpTo(run, expected, chunk_size); + } } INSTANTIATE_TEST_SUITE_P( // @@ -451,15 +505,53 @@ void CheckRleBitPackedDecode(const std::vector& bytes, }); } -/// Run the decode check over a battery of chunk sizes and output offsets. +/// Count both boolean values across the whole `bytes` using CountUpTo, `chunk_size` +/// values at a time, and check against `expected`. +void CheckRleBitPackedCountUpTo(const std::vector& bytes, + const std::vector& expected, + rle_size_t chunk_size) { + ARROW_SCOPED_TRACE("chunk_size = ", chunk_size); + const auto n_vals = static_cast(expected.size()); + + for (bool value : {true, false}) { + ARROW_SCOPED_TRACE("value = ", value); + RleBitPackedToBitmapDecoder decoder(bytes.data(), + static_cast(bytes.size())); + EXPECT_EQ(decoder.exhausted(), n_vals == 0); + + rle_size_t processed = 0; + rle_size_t matching = 0; + while (processed < n_vals) { + const auto want = std::min(chunk_size, n_vals - processed); + const auto result = decoder.CountUpTo(value, want); + EXPECT_EQ(result.processed_count, want) << "at pos " << processed; + ASSERT_GT(result.processed_count, 0) << "at pos " << processed; // break on failure + EXPECT_EQ(result.matching_count, + CountValueSlice(expected, processed, result.processed_count, value)) + << "at pos " << processed; + processed += result.processed_count; + matching += result.matching_count; + } + EXPECT_EQ(processed, n_vals); + EXPECT_TRUE(decoder.exhausted()); + EXPECT_EQ(matching, CountValueSlice(expected, 0, n_vals, value)); + + // Counting past the end processes nothing and leaves the decoder exhausted. + const auto past = decoder.CountUpTo(value, 8); + EXPECT_EQ(past.processed_count, 0); + EXPECT_EQ(past.matching_count, 0); + EXPECT_TRUE(decoder.exhausted()); + } +} + +/// Run the decode and count checks over a battery of chunk sizes and output offsets. void CheckRleBitPackedToBitmap(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 chunk_size : - {rle_size_t{1}, rle_size_t{3}, rle_size_t{7}, rle_size_t{8}, rle_size_t{9}, - rle_size_t{33}, n_vals, n_vals + 1}) { + for (const rle_size_t chunk_size : {1, 3, 7, 8, 9, 33, n_vals, n_vals + 1}) { CheckRleBitPackedDecode(bytes, expected, chunk_size); + CheckRleBitPackedCountUpTo(bytes, expected, chunk_size); // A non-zero output offset forces the first run to start at a non-byte // aligned output position. for (rle_size_t out_offset = 1; out_offset < 8; ++out_offset) { @@ -471,18 +563,35 @@ void CheckRleBitPackedToBitmap(const std::vector& bytes, } // namespace TEST(RleBitPackedToBitmapDecoder, Empty) { - // A default-constructed decoder is already exhausted. + // A default-constructed decoder is already exhausted: it decodes and counts nothing. RleBitPackedToBitmapDecoder decoder; EXPECT_TRUE(decoder.exhausted()); uint8_t out = 0; auto got = decoder.GetBatch(BitmapSpanMut(&out), 8); EXPECT_EQ(got, 0); + for (bool value : {true, false}) { + const auto counted = decoder.CountUpTo(value, 8); + EXPECT_EQ(counted.processed_count, 0); + EXPECT_EQ(counted.matching_count, 0); + } // So is one reset on an empty buffer. decoder.Reset(nullptr, 0); EXPECT_TRUE(decoder.exhausted()); got = decoder.GetBatch(BitmapSpanMut(&out), 8); EXPECT_EQ(got, 0); + + // A zero batch_size is a no-op even with data available. + std::vector bytes; + std::vector expected; + AppendRleRun(bytes, expected, /*value=*/true, /*count=*/8); + decoder.Reset(bytes.data(), static_cast(bytes.size())); + EXPECT_FALSE(decoder.exhausted()); + EXPECT_EQ(decoder.GetBatch(BitmapSpanMut(&out), 0), 0); + const auto counted = decoder.CountUpTo(true, 0); + EXPECT_EQ(counted.processed_count, 0); + EXPECT_EQ(counted.matching_count, 0); + EXPECT_FALSE(decoder.exhausted()); } TEST(RleBitPackedToBitmapDecoder, SingleRleZeros) { From 666caaf642f983da613c145ae0530988dc8f44c3 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 21 Jul 2026 15:26:07 +0200 Subject: [PATCH 24/33] Add BitPackedToBitmapDecoder --- cpp/src/arrow/util/rle_bitmap_internal.h | 46 ++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/cpp/src/arrow/util/rle_bitmap_internal.h b/cpp/src/arrow/util/rle_bitmap_internal.h index f1395a008a9..78954b8af76 100644 --- a/cpp/src/arrow/util/rle_bitmap_internal.h +++ b/cpp/src/arrow/util/rle_bitmap_internal.h @@ -19,6 +19,7 @@ #include #include +#include #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -340,6 +341,51 @@ class RleBitPackedToBitmapDecoder { rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size); }; +/// Minimal decoder for legacy bit packed encoding (BIT_PACKED = 4) into a bitmap. +/// +/// This is the bitmap counterpart of ``BitPackedDecoder``: it decodes a single +/// bit-packed run of booleans (each value stored on a single bit) directly into an +/// Arrow validity bitmap, without the RLE framing. +/// +/// The number of values that the decoder can represent is up to 2^31 - 1. +class BitPackedToBitmapDecoder : private BitPackedRunToBitmapDecoder { + private: + using Base = BitPackedRunToBitmapDecoder; + + public: + using Base::Advance; + using Base::CountUpTo; + using Base::Get; + using Base::GetBatch; + using Base::remaining; + + BitPackedToBitmapDecoder() noexcept = default; + + /// Create a decoder object. + /// + /// @param data and data_size are the raw bytes to decode. + /// @param value_count is the number of values in the run. + BitPackedToBitmapDecoder(const uint8_t* data, rle_size_t data_size, + rle_size_t value_count) noexcept { + Reset(data, data_size, value_count); + } + + void Reset(const uint8_t* data, rle_size_t data_size, rle_size_t value_count) noexcept { + ARROW_DCHECK_GE(value_count, 0); + ARROW_DCHECK_LE(value_count, std::numeric_limits::max()); + const auto run = BitPackedRun{ + /* data= */ data, + /* value_count= */ value_count, + /* value_bit_width= */ 1, + /* max_read_bytes= */ data_size, + }; + return Base::Reset(run); + } + + /// Whether there is still values to iterate over. + bool exhausted() const { return Base::remaining() == 0; } +}; + /************************************************ * RleBitPackedToBitmapDecoder implementation * ************************************************/ From 365ef85f90f148d5065e462b295c67f780362fab Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 21 Jul 2026 15:35:56 +0200 Subject: [PATCH 25/33] Add RleBitPackedToBitmapDecoder::Advance --- cpp/src/arrow/util/rle_bitmap_internal.h | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/cpp/src/arrow/util/rle_bitmap_internal.h b/cpp/src/arrow/util/rle_bitmap_internal.h index 78954b8af76..1f4e61808df 100644 --- a/cpp/src/arrow/util/rle_bitmap_internal.h +++ b/cpp/src/arrow/util/rle_bitmap_internal.h @@ -308,6 +308,10 @@ class RleBitPackedToBitmapDecoder { /// Whether there is still runs to iterate over. 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); + /// 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. @@ -449,6 +453,14 @@ auto RleBitPackedToBitmapDecoder::ProcessValues(Callable&& func, return values_processed; } +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); +} + inline auto RleBitPackedToBitmapDecoder::GetBatch(BitmapSpanMut out, rle_size_t batch_size) -> rle_size_t { return ProcessValues( From b5bc0ab0b2788fe86191fd61833555027e07cc3e Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Tue, 21 Jul 2026 15:36:33 +0200 Subject: [PATCH 26/33] Add LevelToBitmapDecoder --- cpp/src/parquet/column_reader.cc | 159 +++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 2f00a9f6637..61f2a91a419 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -40,6 +40,7 @@ #include "arrow/util/crc32.h" #include "arrow/util/int_util_overflow.h" #include "arrow/util/logging.h" +#include "arrow/util/rle_bitmap_internal.h" #include "arrow/util/rle_encoding_internal.h" #include "arrow/util/unreachable.h" #include "parquet/column_page.h" @@ -225,6 +226,164 @@ auto LevelDecoder::CountUpTo(int16_t value, int32_t batch_size) -> CountUpToResu }; } +/************************** + * LevelToBitmapDecoder * + **************************/ + +/// Decoder for definition levels that writes directly into a validity bitmap. +/// +/// This is the bitmap counterpart of ``LevelDecoder``, specialized for levels +/// encoded on a single bit (a max level of 1), such as the definition levels of a +/// flat, nullable column. Rather than decoding into an ``int16_t`` array and +/// re-encoding into an Arrow validity bitmap, it decodes straight into the bitmap. +/// +/// @see LevelDecoder +class LevelToBitmapDecoder { + public: + using BitmapSpanMut = ::arrow::util::BitmapSpanMut; + using RleBitPackedDecoder = ::arrow::util::RleBitPackedToBitmapDecoder; + using BitPackedDecoder = ::arrow::util::BitPackedToBitmapDecoder; + using CountUpToResult = LevelDecoder::CountUpToResult; + + // TODO we should factor this with the LevelDecoder + + LevelToBitmapDecoder() = default; + + /// Initialize the decoder state with new data from a legacy (V1) page. + /// + /// @return the number of bytes consumed + int32_t SetData(Encoding::type encoding, int16_t max_level, int32_t num_buffered_values, + const uint8_t* data, int32_t data_size); + + /// Initialize the decoder state with new data from a V2 page. + /// + /// Repetition and definition levels in V2 pages are always RLE encoded. + void SetDataV2(int32_t num_bytes, int16_t max_level, int32_t num_buffered_values, + const uint8_t* data); + + /// Decode a batch of levels into `out` and return the number of levels decoded. + int32_t Decode(int32_t batch_size, BitmapSpanMut out); + + /// Advance the decoder and throw away decoded levels. + int32_t Skip(int32_t batch_size); + + /// Advance and count the number of occurrences of `value`. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of elements that were processed. + CountUpToResult CountUpTo(bool value, int32_t batch_size); + + /// Return the max level used in this decoder. + int32_t max_level() const { return max_level_; } + + /// Return the number of values left to be decoded. + int32_t remaining() const { return num_values_remaining_; } + + private: + static void CheckMaxLevel(int16_t max_level); + + std::variant decoder_ = {}; + /// Number of values 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. + int32_t num_values_remaining_ = 0; + int16_t max_level_ = 0; +}; + +void LevelToBitmapDecoder::CheckMaxLevel(int16_t max_level) { + if (ARROW_PREDICT_FALSE(max_level != 1)) { + throw ParquetException( + "LevelToBitmapDecoder only supports levels with a max level of 1."); + } +} + +int32_t LevelToBitmapDecoder::SetData(Encoding::type encoding, int16_t max_level, + int32_t num_buffered_values, const uint8_t* data, + int32_t data_size) { + CheckMaxLevel(max_level); + max_level_ = max_level; + num_values_remaining_ = num_buffered_values; + // Levels with a max level of 1 are encoded on a single bit. + constexpr int32_t value_bit_width = 1; + + switch (encoding) { + case Encoding::RLE: { + if (data_size < 4) { + throw ParquetException("Received invalid levels (corrupt data page?)"); + } + 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?)"); + } + decoder_ = RleBitPackedDecoder( + /* data= */ data + 4, + /* data_size= */ num_bytes); + return 4 + num_bytes; + } + case Encoding::BIT_PACKED: { + int32_t num_bits = 0; + if (MultiplyWithOverflow(num_buffered_values, value_bit_width, &num_bits)) { + throw ParquetException( + "Number of buffered values too large (corrupt data page?)"); + } + 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?)"); + } + // Also passing `value_count` so that the decoder works with zero-width runs. + decoder_ = BitPackedDecoder( + /* data= */ data, + /* data_size= */ num_bytes, + /* value_count= */ num_buffered_values); + return num_bytes; + } + default: + throw ParquetException("Unknown encoding type for levels."); + } + return -1; +} + +void LevelToBitmapDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, + int32_t num_buffered_values, const uint8_t* data) { + CheckMaxLevel(max_level); + if (num_bytes < 0) { + throw ParquetException("Invalid page header (corrupt data page?)"); + } + max_level_ = max_level; + num_values_remaining_ = num_buffered_values; + decoder_ = RleBitPackedDecoder( + /* data= */ data, + /* data_size= */ num_bytes); +} + +int32_t LevelToBitmapDecoder::Decode(int32_t batch_size, BitmapSpanMut out) { + const int32_t num_values = std::min(num_values_remaining_, batch_size); + const int32_t num_decoded = + std::visit([&](auto& dec) { return dec.GetBatch(out, num_values); }, decoder_); + num_values_remaining_ -= num_decoded; + return num_decoded; +} + +int32_t LevelToBitmapDecoder::Skip(int32_t batch_size) { + const int32_t num_values = std::min(num_values_remaining_, batch_size); + const int32_t num_advanced = + std::visit([&](auto& dec) { return dec.Advance(num_values); }, decoder_); + ARROW_DCHECK_EQ(num_values, num_advanced); + num_values_remaining_ -= num_advanced; + return num_advanced; +} + +auto LevelToBitmapDecoder::CountUpTo(bool value, int32_t batch_size) -> CountUpToResult { + const int32_t num_values = std::min(num_values_remaining_, batch_size); + const auto result = + std::visit([&](auto& dec) { return dec.CountUpTo(value, num_values); }, decoder_); + ARROW_DCHECK_EQ(num_values, result.processed_count); + num_values_remaining_ -= result.processed_count; + return { + .matching_count = result.matching_count, + .processed_count = result.processed_count, + }; +} + ReaderProperties default_reader_properties() { static ReaderProperties default_reader_properties; return default_reader_properties; From 6da1c42438fcd7ec1b7302a622062925ad1f72ae Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 22 Jul 2026 10:04:45 +0200 Subject: [PATCH 27/33] Add ValiditySinkBuffer --- cpp/src/parquet/column_reader.cc | 165 ++++++++++++++++++++----------- 1 file changed, 110 insertions(+), 55 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 61f2a91a419..42f6601639e 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -952,11 +952,6 @@ class ValueSinkBuffer : private ValueSinkCursor { void OnNewDictionary(auto& /* decoder */) {} - void delete_back(int64_t count) { - ARROW_DCHECK_LE(count, values_count()); - set_values_count(values_count() - count); - } - [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { const auto decoded = decoder.Decode(write_start(), batch_size); set_values_count(values_count() + batch_size); @@ -1012,6 +1007,85 @@ class ValueSinkBuffer : private ValueSinkCursor { value_type* write_start() { return data() + values_count(); } }; +/************************ + * ValiditySinkBuffer * + ************************/ + +class ValiditySinkBuffer : private ValueSinkCursor { + public: + using ValueSinkCursor::capacity; + using ValueSinkCursor::values_count; + + ValiditySinkBuffer() = default; + + explicit ValiditySinkBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {} + + uint8_t* data() const { return values_->mutable_data_as(); } + + struct ReadResult { + int64_t values_read = 0; + int64_t null_count = 0; + }; + + ReadResult ReadFromLevels(const int16_t* def_levels, int64_t num_def_levels, + const internal::LevelInfo& level_info) { + internal::ValidityBitmapInputOutput validity_io{}; + validity_io.values_read_upper_bound = num_def_levels; + validity_io.valid_bits = data(); + validity_io.valid_bits_offset = values_count(); + DefLevelsToBitmap(def_levels, num_def_levels, level_info, &validity_io); + ARROW_DCHECK_GE(validity_io.values_read, 0); + ARROW_DCHECK_GE(validity_io.null_count, 0); + + // Advance by the number of leaf values (one validity bit each), not by the + // number of definition levels: for repeated/nested columns some def levels + // describe empty or absent lists that produce no leaf value, so + // values_read <= num_def_levels. + set_values_count(values_count() + validity_io.values_read); + return { + .values_read = validity_io.values_read, + .null_count = validity_io.null_count, + }; + } + + std::shared_ptr ReleaseValues(MemoryPool* pool) { + // TODO should we set values_written to zero? + auto result = values_; + const auto byte_count = bytes_for_values(values_count()); + PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true)); + values_ = AllocateBuffer(pool); + reset_capacity(); + return result; + } + + void ReserveValues(int64_t extra_values) { + const auto old_capacity = fit_capacity_for_extra(extra_values); + if (capacity() > old_capacity) { + const auto byte_count = bytes_for_values(capacity()); + PARQUET_THROW_NOT_OK(values_->Resize(byte_count, /*shrink_to_fit=*/false)); + // Zero the newly grown region so that appending bits at a non-byte-aligned + // offset never reads uninitialized memory (avoids Valgrind/MSAN warnings). + const auto old_byte_count = bytes_for_values(old_capacity); + std::memset(data() + old_byte_count, 0, + static_cast(byte_count - old_byte_count)); + } + } + + void ResetValues() { + if (values_count() > 0) { + PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); + ValueSinkCursor::operator=({}); + } + } + + private: + std::shared_ptr<::arrow::ResizableBuffer> values_; + + static int64_t bytes_for_values(int64_t num_values) { + return bit_util::BytesForBits(num_values); + } +}; + /*********************** * ColumnChunkReader * ***********************/ @@ -1689,12 +1763,10 @@ class TypedRecordReader bool read_dense_for_nullable, ValueSink value_sink) : Base(descr, pool, LevelDecoder(descr->max_definition_level()), LevelDecoder(descr->max_repetition_level())), - valid_bits_(AllocateBuffer(pool)), value_sink_(std::move(value_sink)), leaf_info_(leaf_info), def_levels_(AllocateBuffer(pool)), rep_levels_(AllocateBuffer(pool)), - nullable_values_(leaf_info.HasNullableValues()), read_dense_for_nullable_(read_dense_for_nullable) { TypedRecordReader::Reset(); } @@ -1713,7 +1785,9 @@ class TypedRecordReader int64_t null_count() const final { return null_count_; } - bool nullable_values() const final { return nullable_values_; } + /// \brief Indicates if we can have nullable values. Note that repeated fields + /// may or may not be nullable. + bool nullable_values() const final { return leaf_info_.HasNullableValues(); } bool read_dictionary() const final { return kReadDictionary; } @@ -1797,7 +1871,8 @@ class TypedRecordReader value_sink_.OnNewDictionary(decoder); } - void ReadValuesSpaced(int32_t values_with_nulls, int32_t null_count); + void ReadValuesSpaced(int64_t valid_bits_offset, int32_t values_with_nulls, + int32_t null_count); void ReadValuesDense(int32_t values_to_read); @@ -1828,17 +1903,17 @@ class TypedRecordReader void ResetValues(); protected: - /// \brief Each bit corresponds to one element in 'values_' and specifies if it - /// is null or not null. - /// - /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. - std::shared_ptr<::arrow::ResizableBuffer> valid_bits_; - auto value_sink() -> ValueSink& { return value_sink_; } auto value_sink() const -> const ValueSink& { return value_sink_; } private: ValueSink value_sink_; + /// \brief Each bit corresponds to one element in 'values_' and specifies if it + /// is null or not null. + /// + /// Not set if leaf type is not nullable or read_dense_for_nullable_ is true. + // TODO make optional where abscense means read_dense_for_optional + ValiditySinkBuffer valid_bits_; LevelInfo leaf_info_; /// \brief Buffer for definition levels. May contain more levels than @@ -1865,10 +1940,6 @@ class TypedRecordReader int64_t levels_position_ = 0; int64_t levels_capacity_ = 0; - /// \brief Indicates if we can have nullable values. Note that repeated fields - /// may or may not be nullable. - bool nullable_values_; - bool at_record_start_ = true; // If true, we will not leave any space for the null values in the values_ @@ -2161,12 +2232,7 @@ int64_t TypedRecordReader::SkipRecords(int64_t num_records) { template std::shared_ptr TypedRecordReader::ReleaseIsValid() { if (nullable_values()) { - const auto bit_count = bit_util::BytesForBits(values_written()); - auto result = valid_bits_; - PARQUET_THROW_NOT_OK(result->Resize(bit_count, - /*shrink_to_fit=*/true)); - valid_bits_ = AllocateBuffer(this->pool_); - return result; + return valid_bits_.ReleaseValues(this->pool_); } else { return nullptr; } @@ -2264,25 +2330,16 @@ void TypedRecordReader::ReserveLevels(int64_t extra_levels) { template void TypedRecordReader::ReserveIsValid(int64_t extra_values) { if (nullable_values() && !read_dense_for_nullable_) { - const int64_t current_byte_capacity = valid_bits_->size(); - const int64_t bit_capacity = compute_capacity_pow2( - /* capacity= */ 8 * current_byte_capacity, - /* size= */ values_written(), - /* extra_size= */ extra_values); - const int64_t byte_capacity = bit_util::BytesForBits(bit_capacity); - if (current_byte_capacity < byte_capacity) { - PARQUET_THROW_NOT_OK(valid_bits_->Resize(byte_capacity, /*shrink_to_fit=*/false)); - // Avoid valgrind warnings - const int64_t old_valid_bytes = bit_util::BytesForBits(values_written()); - std::memset(valid_bits_->mutable_data() + old_valid_bytes, 0, - static_cast(byte_capacity - old_valid_bytes)); - } + valid_bits_.ReserveValues(extra_values); } } template void TypedRecordReader::Reset() { ResetValues(); + if (!read_dense_for_nullable()) { + valid_bits_ = ValiditySinkBuffer(this->pool_); + } if (levels_written_ > 0) { // Throw away levels from 0 to levels_position_. @@ -2299,12 +2356,13 @@ void TypedRecordReader::SetPageReader(std::unique_ptr } template -void TypedRecordReader::ReadValuesSpaced(int32_t values_with_nulls, +void TypedRecordReader::ReadValuesSpaced(int64_t valid_bits_offset, + int32_t values_with_nulls, int32_t null_count) { const auto decoded = value_sink_.ReadValuesSpaced( *this->current_decoder_.get(), values_with_nulls, null_count, - /* valid_bits= */ valid_bits_->mutable_data(), /* valid_bits_offset= */ - values_written()); + /* valid_bits= */ valid_bits_.data(), + /* valid_bits_offset= */ valid_bits_offset); CheckNumberDecoded(decoded, values_with_nulls); } @@ -2383,20 +2441,17 @@ void TypedRecordReader::ReadSpacedForOptionalOrRepeated( int64_t start_levels_position, int64_t* values_to_read, int64_t* null_count) { // levels_position_ must already be incremented based on number of records // read. - ARROW_DCHECK_GE(levels_position_, start_levels_position); - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = levels_position_ - start_levels_position; - validity_io.valid_bits = valid_bits_->mutable_data(); - validity_io.valid_bits_offset = values_written(); - - DefLevelsToBitmap(def_levels() + start_levels_position, - levels_position_ - start_levels_position, leaf_info_, &validity_io); - *values_to_read = validity_io.values_read - validity_io.null_count; - *null_count = validity_io.null_count; - ARROW_DCHECK_GE(*values_to_read, 0); - ARROW_DCHECK_GE(*null_count, 0); + const int64_t valid_bits_offset = valid_bits_.values_count(); + ARROW_DCHECK_EQ(values_written(), valid_bits_offset); + const auto result = valid_bits_.ReadFromLevels( // + def_levels() + start_levels_position, levels_position_ - start_levels_position, + leaf_info_); + + *values_to_read = result.values_read - result.null_count; + *null_count = result.null_count; + // This is only reading in the current page so this fits in an int32. - ReadValuesSpaced(clamp_to(validity_io.values_read), + ReadValuesSpaced(valid_bits_offset, clamp_to(result.values_read), clamp_to(*null_count)); } @@ -2495,7 +2550,7 @@ template void TypedRecordReader::ResetValues() { if (values_written() > 0) { value_sink_.ResetValues(); - PARQUET_THROW_NOT_OK(valid_bits_->Resize(0, /*shrink_to_fit=*/false)); + valid_bits_.ResetValues(); null_count_ = 0; } } From 6e9ff4d1730d145b490fe6dca2c8a7e57f80d83c Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 22 Jul 2026 10:29:48 +0200 Subject: [PATCH 28/33] Auto reserve space in sinks: --- cpp/src/parquet/column_reader.cc | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 42f6601639e..3de2b785f37 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -953,6 +953,7 @@ class ValueSinkBuffer : private ValueSinkCursor { void OnNewDictionary(auto& /* decoder */) {} [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { + ReserveValues(batch_size); const auto decoded = decoder.Decode(write_start(), batch_size); set_values_count(values_count() + batch_size); return decoded; @@ -961,6 +962,7 @@ class ValueSinkBuffer : private ValueSinkCursor { [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, int32_t null_count, const uint8_t* valid_bits, int64_t valid_bits_offset) { + ReserveValues(batch_size); const auto decoded = decoder.DecodeSpaced(write_start(), batch_size, null_count, valid_bits, valid_bits_offset); set_values_count(values_count() + batch_size); @@ -1029,6 +1031,8 @@ class ValiditySinkBuffer : private ValueSinkCursor { ReadResult ReadFromLevels(const int16_t* def_levels, int64_t num_def_levels, const internal::LevelInfo& level_info) { + // At most one validity bit is written per definition level. + ReserveValues(num_def_levels); internal::ValidityBitmapInputOutput validity_io{}; validity_io.values_read_upper_bound = num_def_levels; validity_io.valid_bits = data(); @@ -2457,12 +2461,8 @@ void TypedRecordReader::ReadSpacedForOptionalOrRepeated( template int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { - // Conservative upper bound - const int64_t possible_num_values = - std::max(num_records, levels_written_ - levels_position_); - ReserveValues(possible_num_values); - ReserveIsValid(possible_num_values); - + // The value and validity sinks reserve their own capacity as they read, so + // there is no need to pre-reserve an upper bound here. const int64_t start_levels_position = levels_position_; // To be updated by the function calls below for each of the repetition @@ -2688,10 +2688,7 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records return 0; } - Reserve(num_records); - int64_t records_read = 0; - do { // Is there more data to read in this row group? if (!this->ProcessToMoreData()) { @@ -2777,6 +2774,7 @@ class ArrayValuesSink : private ValueSinkCursor { void OnNewDictionary(auto& /* decoder */) {} [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) { + ReserveValues(batch_size); // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = decoder.DecodeArrowNonNull(batch_size, &builder_); mark_values_as_written(batch_size); @@ -2786,6 +2784,7 @@ class ArrayValuesSink : private ValueSinkCursor { [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size, int32_t null_count, const uint8_t* valid_bits, int64_t valid_bits_offset) { + ReserveValues(batch_size); // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = decoder.DecodeArrow(batch_size, null_count, valid_bits, valid_bits_offset, &builder_); @@ -2999,6 +2998,7 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { void ResetValues() { ValueSinkCursor::operator=({}); } [[nodiscard]] int64_t ReadValuesDense(auto& decoder, int32_t batch_size) { + ReserveValues(batch_size); // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = [&]() -> int64_t { if (auto* dict_decoder = dynamic_cast(&decoder)) { @@ -3013,6 +3013,7 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { [[nodiscard]] int64_t ReadValuesSpaced(auto& decoder, int32_t batch_size, int32_t null_count, const uint8_t* valid_bits, int64_t valid_bits_offset) { + ReserveValues(batch_size); // TODO once we have a validity sink: we can reset the validity to save some space const int64_t decoded = [&]() -> int64_t { if (auto* dict_decoder = dynamic_cast(&decoder)) { From a5518c741157908ff5180da356c02914abf459ae Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 22 Jul 2026 12:35:43 +0200 Subject: [PATCH 29/33] Fix integer types --- cpp/src/parquet/column_reader.cc | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 3de2b785f37..512777083a7 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1458,16 +1458,15 @@ int64_t ColumnChunkReader::Skip(int64_t num_values_to_skip) { while (values_to_skip > 0 && ProcessToMoreData()) { // If the number of values to skip is more than the number of undecoded values, skip // the whole Page without decoding levels or values. - const int64_t available_values = this->available_values_current_page(); + const int32_t available_values = this->available_values_current_page(); if (values_to_skip >= available_values) { values_to_skip -= available_values; this->ConsumeBufferedValues(available_values); } else { - // 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); + // The whole batch fits inside this Page (`<= availables_values`) + const auto batch_size = static_cast(values_to_skip); - int64_t non_null = AdvanceLevels(batch_size); + const int32_t non_null = AdvanceLevels(batch_size); // Skip the corresponding data values. this->current_decoder_.Skip(non_null); From f5ec507aa082819a2f2b5bb032c4223d4f5e5343 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Wed, 22 Jul 2026 15:29:53 +0200 Subject: [PATCH 30/33] Nit in ColumnChunkReader --- cpp/src/parquet/column_reader.cc | 109 +++++++++++++++++++------------ 1 file changed, 68 insertions(+), 41 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 512777083a7..5fc8094a6ba 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -91,13 +91,33 @@ inline void CheckNumberDecoded(int64_t number_decoded, int64_t expected) { constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = "Number of decoded rep / def levels do not match num_values in page header"; +/// True if a T can hold a V. +template +inline constexpr bool can_hold_v = std::in_range(std::numeric_limits::min()) && + std::in_range(std::numeric_limits::max()); + +/// Clamp to a given type. template constexpr T clamp_to(U val) { + if constexpr (can_hold_v) { + return static_cast(val); + } constexpr U kMax = std::numeric_limits::max(); constexpr U kMin = std::numeric_limits::min(); return static_cast(std::clamp(val, kMin, kMax)); } +/// Return the int that can be represented by the other. +template +using smallest_int_t = std::conditional_t, U, T>; + +/// Return the min in the smallest integer size of its inputs. +template +constexpr auto narrow_min(T x, U y) -> smallest_int_t { + using V = smallest_int_t; + return std::min(clamp_to(x), clamp_to(y)); +} + } // namespace /****************** @@ -1181,9 +1201,11 @@ class ColumnChunkReader { rep_levels_decoder_(std::move(rep_levels_decoder)) {} void SetPageReader(std::unique_ptr reader) { + current_page_ = nullptr; pager_ = std::move(reader); - // Dictionary decoders must be reset when advancing row groups + current_decoder_.SetDecoder(nullptr); decoders_.clear(); + current_encoding_ = Encoding::UNKNOWN; } bool HasPageReader() const { return pager_ != nullptr; } @@ -1197,14 +1219,6 @@ class ColumnChunkReader { /// Check the encoding of the current page or throw an exception. void CheckEncodingIs(Encoding::type encoding); - private: - // Declared here because `decoders_` below refers to it. - using DecoderType = TypedDecoder; - - Derived& derived() { return static_cast(*this); } - const Derived& derived() const { return static_cast(*this); } - - protected: int32_t ReadDefinitionLevels(int32_t batch_size, int16_t* levels) { if (max_def_level() == 0) { return 0; @@ -1219,10 +1233,6 @@ class ColumnChunkReader { return rep_levels_decoder_.Decode(batch_size, levels); } - const ColumnDescriptor* descr_; - ::arrow::MemoryPool* pool_; - SkippableTypedDecoder current_decoder_; - // Available values in the current data page, value includes repeated values and nulls. int32_t available_values_current_page() const { const int32_t out = num_buffered_values_ - num_decoded_values_; @@ -1238,7 +1248,14 @@ class ColumnChunkReader { int64_t Skip(int64_t num_values_to_skip); + protected: + const ColumnDescriptor* descr_; + ::arrow::MemoryPool* pool_; + SkippableTypedDecoder current_decoder_; + private: + using DecoderType = TypedDecoder; + LvlDec def_levels_decoder_; LvlDec rep_levels_decoder_; // Map of encoding type to the respective decoder object. For example, a @@ -1246,7 +1263,7 @@ class ColumnChunkReader { // plain-encoded data. std::unordered_map> decoders_; std::unique_ptr pager_; - std::shared_ptr current_page_; + std::shared_ptr current_page_; // The total number of values stored in the data page. This is the maximum of // the number of encoded definition levels or encoded values. For // non-repeated, required columns, this is equal to the number of encoded @@ -1259,11 +1276,17 @@ class ColumnChunkReader { int32_t num_decoded_values_ = 0; Encoding::type current_encoding_ = Encoding::UNKNOWN; + Derived& derived() { return static_cast(*this); } + const Derived& derived() const { return static_cast(*this); } + // Advance to the next data page bool ReadNewPage(); void ConfigureDictionary(const DictionaryPage* page); + template + void InitializeDataPage(std::shared_ptr page); + // Get a decoder object for this page or create a new decoder if this is the // first page with this encoding. void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size); @@ -1305,31 +1328,20 @@ template bool ColumnChunkReader::ReadNewPage() { // Loop until we find the next data page. while (true) { - current_page_ = pager_->NextPage(); - if (!current_page_) { + std::shared_ptr page = pager_->NextPage(); + if (!page) { // EOS return false; } - if (current_page_->type() == PageType::DICTIONARY_PAGE) { - ConfigureDictionary(static_cast(current_page_.get())); + if (page->type() == PageType::DICTIONARY_PAGE) { + ConfigureDictionary(static_cast(page.get())); continue; - } else if (current_page_->type() == PageType::DATA_PAGE) { - const auto& page = static_cast(*current_page_); - const int64_t levels_byte_size = - InitializeV1Levels(page, def_levels_decoder_, rep_levels_decoder_); - num_buffered_values_ = page.num_values(); - num_decoded_values_ = 0; - InitializeDataDecoder(page, levels_byte_size); + } else if (page->type() == PageType::DATA_PAGE) { + InitializeDataPage(std::static_pointer_cast(page)); return true; - } else if (current_page_->type() == PageType::DATA_PAGE_V2) { - const auto& page = static_cast(*current_page_); - const int64_t levels_byte_size = - InitializeV2Levels(page, def_levels_decoder_, rep_levels_decoder_); - num_buffered_values_ = page.num_values(); - num_buffered_values_ = page.num_values(); - num_decoded_values_ = 0; - InitializeDataDecoder(page, levels_byte_size); + } else if (page->type() == PageType::DATA_PAGE_V2) { + InitializeDataPage(std::static_pointer_cast(page)); return true; } else { // We don't know what this page type is. We're allowed to skip non-data @@ -1377,6 +1389,24 @@ void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) ARROW_DCHECK(current_decoder_); } +template +template +void ColumnChunkReader::InitializeDataPage(std::shared_ptr page) { + ARROW_DCHECK_NE(page, nullptr); + current_page_ = std::static_pointer_cast(page); + int64_t byte_size = 0; + if constexpr (std::is_same_v) { + byte_size = InitializeV1Levels(*page, def_levels_decoder_, rep_levels_decoder_); + } else if constexpr (std::is_same_v) { + byte_size = InitializeV2Levels(*page, def_levels_decoder_, rep_levels_decoder_); + } else { + static_assert(false, "Unknown data page type"); + } + num_buffered_values_ = page->num_values(); + num_decoded_values_ = 0; + InitializeDataDecoder(*page, byte_size); +} + template void ColumnChunkReader::InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size) { @@ -2001,10 +2031,8 @@ int64_t TypedRecordReader::ReadRecords(int64_t num_records) { /// We perform multiple batch reads until we either exhaust the row group /// or observe the desired number of records - const int64_t batch_size_64 = - std::min(level_batch_size, this->available_values_current_page()); - // available_values_current_page fits in int32_t - const auto batch_size = static_cast(batch_size_64); + const int32_t batch_size = + narrow_min(level_batch_size, this->available_values_current_page()); // No more data in column if (batch_size == 0) { @@ -2032,8 +2060,8 @@ int64_t TypedRecordReader::ReadRecords(int64_t num_records) { records_read += ReadRecordData(num_records - records_read); } else { // No repetition and definition levels, we can read values directly - const auto batch_size = std::min(num_records - records_read, batch_size_64); - records_read += ReadRecordData(batch_size); + const auto count = narrow_min(num_records - records_read, batch_size); + records_read += ReadRecordData(count); } } @@ -2695,8 +2723,7 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records } const int32_t batch_size = - std::min(clamp_to(num_records - records_read), - this->available_values_current_page()); + narrow_min(num_records - records_read, this->available_values_current_page()); ReadValuesDense(batch_size); this->ConsumeBufferedValues(batch_size); From d74d4ad3ca19e4616196ca52ce771bfac8c231ae Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 23 Jul 2026 17:54:34 +0200 Subject: [PATCH 31/33] Eagerly read dict --- cpp/src/parquet/column_reader.cc | 179 +++++++++++++------------------ 1 file changed, 77 insertions(+), 102 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 5fc8094a6ba..a71d8d394d3 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -1180,16 +1180,17 @@ int64_t InitializeV2Levels(const DataPageV2& page, LvlDec& def_dec, LvlDec& rep_ return total_levels_length; } -/// Traits of a concrete column chunk reader, for use by `ColumnChunkReader`. -template -struct reader_trait; - /// Read through the multiple pages of a column chunk. -template +/// +/// `Traits` must provide two member types: +/// - `DType`: the physical Parquet type of the column; +/// - `level_decoder`: the decoder type used for definition/repetition levels. +/// Each concrete reader defines its own trait struct. +template class ColumnChunkReader { public: - using DType = typename reader_trait::DType; - using LvlDec = typename reader_trait::level_decoder; + using DType = typename Traits::DType; + using LvlDec = typename Traits::level_decoder; using value_type = typename DType::c_type; ColumnChunkReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, @@ -1214,7 +1215,7 @@ class ColumnChunkReader { /// /// If the current page is exhausted, it will process more pages until some data /// page is found. - bool ProcessToMoreData(); + bool EnsureDataPage(); /// Check the encoding of the current page or throw an exception. void CheckEncodingIs(Encoding::type encoding); @@ -1244,7 +1245,7 @@ class ColumnChunkReader { int16_t max_rep_level() const; - void ConsumeBufferedValues(int64_t num_values) { num_decoded_values_ += num_values; } + void MarkValuesAsConsumed(int64_t num_values) { num_decoded_values_ += num_values; } int64_t Skip(int64_t num_values_to_skip); @@ -1276,9 +1277,6 @@ class ColumnChunkReader { int32_t num_decoded_values_ = 0; Encoding::type current_encoding_ = Encoding::UNKNOWN; - Derived& derived() { return static_cast(*this); } - const Derived& derived() const { return static_cast(*this); } - // Advance to the next data page bool ReadNewPage(); @@ -1302,8 +1300,8 @@ class ColumnChunkReader { * ColumnChunkReader Implementation * **************************************/ -template -void ColumnChunkReader::CheckEncodingIs(Encoding::type encoding) { +template +void ColumnChunkReader::CheckEncodingIs(Encoding::type encoding) { if (current_encoding_ != encoding) { auto msg = std::format("Unexpected data page encoding. Expected {}, got {}", @@ -1312,8 +1310,8 @@ void ColumnChunkReader::CheckEncodingIs(Encoding::type encoding) { } } -template -bool ColumnChunkReader::ProcessToMoreData() { +template +bool ColumnChunkReader::EnsureDataPage() { // Either there is no data page available yet, or the data page has been // exhausted if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { @@ -1324,8 +1322,8 @@ bool ColumnChunkReader::ProcessToMoreData() { return true; } -template -bool ColumnChunkReader::ReadNewPage() { +template +bool ColumnChunkReader::ReadNewPage() { // Loop until we find the next data page. while (true) { std::shared_ptr page = pager_->NextPage(); @@ -1352,8 +1350,8 @@ bool ColumnChunkReader::ReadNewPage() { return true; } -template -void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) { +template +void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) { int encoding = static_cast(page->encoding()); if (page->encoding() == Encoding::PLAIN_DICTIONARY || page->encoding() == Encoding::PLAIN) { @@ -1378,7 +1376,6 @@ void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) std::unique_ptr> decoder = MakeDictDecoder(descr_, pool_); decoder->SetDict(dictionary.get()); - derived().OnNewDictionary(*decoder); decoders_[encoding] = std::unique_ptr(dynamic_cast(decoder.release())); } else { @@ -1389,9 +1386,9 @@ void ColumnChunkReader::ConfigureDictionary(const DictionaryPage* page) ARROW_DCHECK(current_decoder_); } -template +template template -void ColumnChunkReader::InitializeDataPage(std::shared_ptr page) { +void ColumnChunkReader::InitializeDataPage(std::shared_ptr page) { ARROW_DCHECK_NE(page, nullptr); current_page_ = std::static_pointer_cast(page); int64_t byte_size = 0; @@ -1407,9 +1404,9 @@ void ColumnChunkReader::InitializeDataPage(std::shared_ptr page) { InitializeDataDecoder(*page, byte_size); } -template -void ColumnChunkReader::InitializeDataDecoder(const DataPage& page, - int64_t levels_byte_size) { +template +void ColumnChunkReader::InitializeDataDecoder(const DataPage& page, + int64_t levels_byte_size) { const uint8_t* buffer = page.data() + levels_byte_size; const int64_t data_size = page.size() - levels_byte_size; @@ -1454,18 +1451,18 @@ void ColumnChunkReader::InitializeDataDecoder(const DataPage& page, static_cast(data_size)); } -template -int16_t ColumnChunkReader::max_def_level() const { +template +int16_t ColumnChunkReader::max_def_level() const { return def_levels_decoder_.max_level(); } -template -int16_t ColumnChunkReader::max_rep_level() const { +template +int16_t ColumnChunkReader::max_rep_level() const { return rep_levels_decoder_.max_level(); } -template -int32_t ColumnChunkReader::AdvanceLevels(int32_t num_levels) { +template +int32_t ColumnChunkReader::AdvanceLevels(int32_t num_levels) { int max_count = num_levels; // Advance the definition levels, counting how many correspond to present // (non-null) values that must be skipped in the data decoder. @@ -1481,17 +1478,17 @@ int32_t ColumnChunkReader::AdvanceLevels(int32_t num_levels) { return max_count; } -template -int64_t ColumnChunkReader::Skip(int64_t num_values_to_skip) { +template +int64_t ColumnChunkReader::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 && ProcessToMoreData()) { + while (values_to_skip > 0 && EnsureDataPage()) { // If the number of values to skip is more than the number of undecoded values, skip // the whole Page without decoding levels or values. const int32_t available_values = this->available_values_current_page(); if (values_to_skip >= available_values) { values_to_skip -= available_values; - this->ConsumeBufferedValues(available_values); + this->MarkValuesAsConsumed(available_values); } else { // The whole batch fits inside this Page (`<= availables_values`) const auto batch_size = static_cast(values_to_skip); @@ -1500,7 +1497,7 @@ int64_t ColumnChunkReader::Skip(int64_t num_values_to_skip) { // Skip the corresponding data values. this->current_decoder_.Skip(non_null); - this->ConsumeBufferedValues(batch_size); + this->MarkValuesAsConsumed(batch_size); values_to_skip -= batch_size; } } @@ -1511,39 +1508,34 @@ int64_t ColumnChunkReader::Skip(int64_t num_values_to_skip) { * TypedColumnReaderImpl * ***************************/ -template -class TypedColumnReaderImpl; - template -struct reader_trait> { +struct TypedColumnReaderImplTraits { using DType = D; using level_decoder = LevelDecoder; }; template -class TypedColumnReaderImpl : public TypedColumnReader, - public ColumnChunkReader> { +class TypedColumnReaderImpl + : public TypedColumnReader, + public ColumnChunkReader> { public: using T = typename DType::c_type; + using Base = ColumnChunkReader>; TypedColumnReaderImpl(const ColumnDescriptor* descr, std::unique_ptr pager, ::arrow::MemoryPool* pool) - : ColumnChunkReader>( - descr, pool, LevelDecoder(descr->max_definition_level()), - LevelDecoder(descr->max_repetition_level())) { + : Base(descr, pool, LevelDecoder(descr->max_definition_level()), + LevelDecoder(descr->max_repetition_level())) { this->SetPageReader(std::move(pager)); } - bool HasNext() override { return this->ProcessToMoreData(); } - - /// Called by ColumnChunkReader. - void OnNewDictionary(DictDecoder& /* decoder */) {} + bool HasNext() override { return this->EnsureDataPage(); } int64_t ReadBatch(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, T* values, int64_t* values_read) override; int64_t Skip(int64_t num_values_to_skip) override { - return ColumnChunkReader>::Skip(num_values_to_skip); + return Base::Skip(num_values_to_skip); } Type::type type() const override { return this->descr_->physical_type(); } @@ -1660,7 +1652,7 @@ int64_t TypedColumnReaderImpl::ReadBatchWithDictionary( ss << "Read 0 values, expected " << expected_values; ParquetException::EofException(ss.str()); } - this->ConsumeBufferedValues(total_indices); + this->MarkValuesAsConsumed(total_indices); return total_indices; } @@ -1706,7 +1698,7 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size_64, ss << "Read 0 values, expected " << expected_values; ParquetException::EofException(ss.str()); } - this->ConsumeBufferedValues(total_values); + this->MarkValuesAsConsumed(total_values); return total_values; } @@ -1769,28 +1761,18 @@ concept can_cout = requires(std::ostream& os, const T& value) { template class TypedRecordReader; -// `reader_trait` can only be specialized in the namespace it is declared in. -} // namespace -} // namespace internal - -namespace { -template -struct reader_trait> { +template +struct TypedRecordReaderTraits { using DType = D; using level_decoder = LevelDecoder; }; -} // namespace - -namespace internal { -namespace { template -class TypedRecordReader - : public ColumnChunkReader>, - virtual public RecordReader { +class TypedRecordReader : public ColumnChunkReader>, + virtual public RecordReader { public: using T = typename DType::c_type; - using Base = ColumnChunkReader>; + using Base = ColumnChunkReader>; TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable, ValueSink value_sink) @@ -1900,10 +1882,6 @@ class TypedRecordReader const ColumnDescriptor* descr() const override { return this->descr_; } - void OnNewDictionary(DictDecoder& decoder) { - value_sink_.OnNewDictionary(decoder); - } - void ReadValuesSpaced(int64_t valid_bits_offset, int32_t values_with_nulls, int32_t null_count); @@ -1989,7 +1967,7 @@ class TypedRecordReader template const void* TypedRecordReader::ReadDictionary(int32_t* dictionary_length) { - if (!this->current_decoder_ && !this->ProcessToMoreData()) { + if (!this->current_decoder_ && !this->EnsureDataPage()) { *dictionary_length = 0; return nullptr; } @@ -2018,7 +1996,7 @@ int64_t TypedRecordReader::ReadRecords(int64_t num_records) { // enough records while (!at_record_start_ || records_read < num_records) { // Is there more data to read in this row group? - if (!this->ProcessToMoreData()) { + if (!this->EnsureDataPage()) { if (!at_record_start_) { // We ended the row group while inside a record that we haven't seen // the end of yet. So increment the record count for the last record in @@ -2129,7 +2107,7 @@ int64_t TypedRecordReader::SkipRecordsInBufferNonRepeated( ReadAndThrowAwayValues(values_to_read); // Mark the levels as read in the underlying column reader. - this->ConsumeBufferedValues(skipped_records); + this->MarkValuesAsConsumed(skipped_records); return skipped_records; } @@ -2149,7 +2127,7 @@ int64_t TypedRecordReader::DelimitAndSkipRecordsInBuffer( // Mark those levels and values as consumed in the underlying page. // This must be done before we throw away levels since it updates // levels_position_ and levels_written_. - this->ConsumeBufferedValues(levels_position_ - start_levels_position); + this->MarkValuesAsConsumed(levels_position_ - start_levels_position); // Updated levels_position_ and levels_written_. ThrowAwayLevels(start_levels_position); return skipped_records; @@ -2175,7 +2153,7 @@ int64_t TypedRecordReader::SkipRecordsRepeated(int64_t num_records while (!at_record_start_ || skipped_records < num_records) { // Is there more data to read in this row group? // HasNextInternal() will advance to the next page if necessary. - if (!this->ProcessToMoreData()) { + if (!this->EnsureDataPage()) { if (!at_record_start_) { // We ended the row group while inside a record that we haven't seen // the end of yet. So increment the record count for the last record @@ -2384,6 +2362,12 @@ template void TypedRecordReader::SetPageReader(std::unique_ptr reader) { at_record_start_ = true; Base::SetPageReader(std::move(reader)); + // At most one dictionary in Parquet column chunk and it has to be the first page. + if (this->HasPageReader() && this->EnsureDataPage()) { + if (auto* dict = dynamic_cast*>(this->current_decoder_.get())) { + value_sink_.OnNewDictionary(*dict); + } + } } template @@ -2528,10 +2512,10 @@ int64_t TypedRecordReader::ReadRecordData(int64_t num_records) { // Total values, including null spaces, if any if (this->max_def_level() > 0) { // Optional, repeated, or some mix thereof - this->ConsumeBufferedValues(levels_position_ - start_levels_position); + this->MarkValuesAsConsumed(levels_position_ - start_levels_position); } else { // Flat, non-repeated - this->ConsumeBufferedValues(values_to_read); + this->MarkValuesAsConsumed(values_to_read); } return records_read; @@ -2590,31 +2574,20 @@ void TypedRecordReader::ResetValues() { template class RequiredTypedRecordReader; -// `reader_trait` can only be specialized in the namespace it is declared in. -} // namespace -} // namespace internal - -namespace { -template -struct reader_trait> { +template +struct RequiredTypedRecordReaderTraits { using DType = D; using level_decoder = LevelDecoder; }; -} // namespace - -namespace internal { -namespace { template , bool kReadDictionary = false> class RequiredTypedRecordReader - : public ColumnChunkReader< - RequiredTypedRecordReader>, + : public ColumnChunkReader>, virtual public RecordReader { public: using T = typename DType::c_type; - using Base = - ColumnChunkReader>; + using Base = ColumnChunkReader>; RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool, ValueSink value_sink) @@ -2663,7 +2636,13 @@ class RequiredTypedRecordReader void Reset() final; void SetPageReader(std::unique_ptr reader) final { - return Base::SetPageReader(std::move(reader)); + Base::SetPageReader(std::move(reader)); + // At most one dictionary in Parquet column chunk and it has to be the first page. + if (this->HasPageReader() && this->EnsureDataPage()) { + if (auto* dict = dynamic_cast*>(this->current_decoder_.get())) { + value_sink_.OnNewDictionary(*dict); + } + } } bool HasMoreData() const override { @@ -2672,10 +2651,6 @@ class RequiredTypedRecordReader const ColumnDescriptor* descr() const final { return this->descr_; } - void OnNewDictionary(DictDecoder& decoder) { - value_sink_.OnNewDictionary(decoder); - } - void DebugPrintState() final; protected: @@ -2697,7 +2672,7 @@ class RequiredTypedRecordReader template const void* RequiredTypedRecordReader::ReadDictionary( int32_t* dictionary_length) { - if (!this->current_decoder_ && !this->ProcessToMoreData()) { + if (!this->current_decoder_ && !this->EnsureDataPage()) { *dictionary_length = 0; return nullptr; } @@ -2718,14 +2693,14 @@ int64_t RequiredTypedRecordReader::ReadRecords(int64_t num_records int64_t records_read = 0; do { // Is there more data to read in this row group? - if (!this->ProcessToMoreData()) { + if (!this->EnsureDataPage()) { break; } const int32_t batch_size = narrow_min(num_records - records_read, this->available_values_current_page()); ReadValuesDense(batch_size); - this->ConsumeBufferedValues(batch_size); + this->MarkValuesAsConsumed(batch_size); records_read += batch_size; } while (records_read < num_records); From a99e1313b87dd09f056d946556b4a9c45b8f4f1d Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 24 Jul 2026 16:07:59 +0200 Subject: [PATCH 32/33] Cleanups --- cpp/src/parquet/column_reader.cc | 79 ++++++++++++++++---------------- 1 file changed, 39 insertions(+), 40 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index a71d8d394d3..6bc9fdbd018 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -118,6 +119,12 @@ constexpr auto narrow_min(T x, U y) -> smallest_int_t { return std::min(clamp_to(x), clamp_to(y)); } +/// Whether values of type T can be printed to `std::cout`. +template +concept can_cout = requires(std::ostream& os, const T& value) { + os << value; +}; // NOLINT(readability/braces) + } // namespace /****************** @@ -294,7 +301,7 @@ class LevelToBitmapDecoder { CountUpToResult CountUpTo(bool value, int32_t batch_size); /// Return the max level used in this decoder. - int32_t max_level() const { return max_level_; } + int32_t max_level() const { return 1; } /// Return the number of values left to be decoded. int32_t remaining() const { return num_values_remaining_; } @@ -306,7 +313,6 @@ class LevelToBitmapDecoder { /// Number of values 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. int32_t num_values_remaining_ = 0; - int16_t max_level_ = 0; }; void LevelToBitmapDecoder::CheckMaxLevel(int16_t max_level) { @@ -320,7 +326,6 @@ int32_t LevelToBitmapDecoder::SetData(Encoding::type encoding, int16_t max_level int32_t num_buffered_values, const uint8_t* data, int32_t data_size) { CheckMaxLevel(max_level); - max_level_ = max_level; num_values_remaining_ = num_buffered_values; // Levels with a max level of 1 are encoded on a single bit. constexpr int32_t value_bit_width = 1; @@ -368,7 +373,6 @@ void LevelToBitmapDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, if (num_bytes < 0) { throw ParquetException("Invalid page header (corrupt data page?)"); } - max_level_ = max_level; num_values_remaining_ = num_buffered_values; decoder_ = RleBitPackedDecoder( /* data= */ data, @@ -1014,6 +1018,17 @@ class ValueSinkBuffer : private ValueSinkCursor { } } + void DebugPrintState() { + const T* vals = data(); + for (int64_t i = 0; i < values_count(); ++i) { + if constexpr (can_cout) { + std::cout << vals[i] << ' '; + } else { + std::cout << "? "; + } + } + } + private: std::shared_ptr<::arrow::ResizableBuffer> values_; @@ -1758,9 +1773,6 @@ concept can_cout = requires(std::ostream& os, const T& value) { * TypedRecordReader * ***********************/ -template -class TypedRecordReader; - template struct TypedRecordReaderTraits { using DType = D; @@ -2544,16 +2556,7 @@ void TypedRecordReader::DebugPrintState() { } std::cout << "values: "; - if constexpr (can_cout) { - const T* vals = reinterpret_cast(this->values()); - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << vals[i] << " "; - } - } else { - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << "? "; - } - } + value_sink_.DebugPrintState(); std::cout << std::endl; } @@ -2570,16 +2573,13 @@ void TypedRecordReader::ResetValues() { * RequiredTypedRecordReader * *******************************/ -// TODO can we reduce some code share with TypedRecordREader ? -template -class RequiredTypedRecordReader; - template struct RequiredTypedRecordReaderTraits { using DType = D; using level_decoder = LevelDecoder; }; +// TODO can we reduce some code share with TypedRecordREader ? template , bool kReadDictionary = false> class RequiredTypedRecordReader @@ -2725,16 +2725,7 @@ void RequiredTypedRecordReader::Reset() { template void RequiredTypedRecordReader::DebugPrintState() { std::cout << "values: "; - if constexpr (can_cout) { - const T* vals = reinterpret_cast(this->values()); - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << vals[i] << " "; - } - } else { - for (int64_t i = 0; i < this->values_written(); ++i) { - std::cout << "? "; - } - } + value_sink_.DebugPrintState(); std::cout << std::endl; } @@ -2799,6 +2790,8 @@ class ArrayValuesSink : private ValueSinkCursor { auto builder() -> BuilderType& { return builder_; } auto builder() const -> const BuilderType& { return builder_; } + void DebugPrintState() { std::cout << ""; } + private: BuilderType builder_; }; @@ -2970,12 +2963,16 @@ class ByteArrayChunkedRecordReader final auto accumulator() -> Builder& { return this->value_sink().builder(); } }; +/***************************** + * ValuesSinkByteArrayDict * + *****************************/ + /// Decodes byte array values into a ::arrow::BinaryDictionary32Builder. /// /// If the current decoder is dictionary encoded, the dictionary indices are /// decoded directly, otherwise the values are decoded and looked up in the /// builder's memo table. -class ByteArrayDictionaryValuesSink : private ValueSinkCursor { +class ValuesSinkByteArrayDict : private ValueSinkCursor { public: using value_type = ByteArray; using Builder = ::arrow::BinaryDictionary32Builder; @@ -2983,7 +2980,7 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { using ValueSinkCursor::capacity; using ValueSinkCursor::values_count; - explicit ByteArrayDictionaryValuesSink(::arrow::MemoryPool* pool) + explicit ValuesSinkByteArrayDict(::arrow::MemoryPool* pool) : builder_(std::make_unique(pool)) {} value_type* data() const { return nullptr; } @@ -3038,6 +3035,8 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { decoder.InsertDictionary(builder_.get()); } + void DebugPrintState() { std::cout << ""; } + /// Finish the builder and return all the chunks accumulated so far. std::shared_ptr<::arrow::ChunkedArray> FlushChunks() { FlushBuilder(); @@ -3064,10 +3063,14 @@ class ByteArrayDictionaryValuesSink : private ValueSinkCursor { } }; +/************************************* + * ByteArrayDictionaryRecordReader * + *************************************/ + template struct byte_array_dictionary_record_reader { using DType = ByteArrayType; - using ValueSink = ByteArrayDictionaryValuesSink; + using ValueSink = ValuesSinkByteArrayDict; using type = record_reader_base_t; }; @@ -3141,9 +3144,6 @@ std::shared_ptr MakeByteArrayRecordReader( } } -} // namespace - -namespace { template std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, @@ -3205,9 +3205,8 @@ std::shared_ptr RecordReader::Make( read_dense_for_nullable); default: { // PARQUET-1481: This can occur if the file is corrupt - std::stringstream ss; - ss << "Invalid physical column type: " << static_cast(descr->physical_type()); - throw ParquetException(ss.str()); + const auto type = static_cast(descr->physical_type()); + throw ParquetException(std::format("Invalid physical column type: {}", type)); } } // Unreachable code, but suppress compiler warning From cf41de51b30b3ed61bd4bd9e91da106e152e37ff Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 24 Jul 2026 17:00:44 +0200 Subject: [PATCH 33/33] Add FlatOptionalTypedRecordReader --- cpp/src/parquet/column_reader.cc | 314 +++++++++++++++++++++++++++---- 1 file changed, 274 insertions(+), 40 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 6bc9fdbd018..7421ae1ea39 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -36,6 +36,7 @@ #include "arrow/chunked_array.h" #include "arrow/type.h" #include "arrow/util/bit_util.h" +#include "arrow/util/bitmap_ops.h" #include "arrow/util/checked_cast.h" #include "arrow/util/compression.h" #include "arrow/util/crc32.h" @@ -1059,13 +1060,14 @@ class ValiditySinkBuffer : private ValueSinkCursor { uint8_t* data() const { return values_->mutable_data_as(); } + template struct ReadResult { - int64_t values_read = 0; - int64_t null_count = 0; + Int values_read = 0; + Int null_count = 0; }; - ReadResult ReadFromLevels(const int16_t* def_levels, int64_t num_def_levels, - const internal::LevelInfo& level_info) { + ReadResult ReadFromLevels(const int16_t* def_levels, int64_t num_def_levels, + const internal::LevelInfo& level_info) { // At most one validity bit is written per definition level. ReserveValues(num_def_levels); internal::ValidityBitmapInputOutput validity_io{}; @@ -1087,6 +1089,24 @@ class ValiditySinkBuffer : private ValueSinkCursor { }; } + ReadResult ReadFromDecoder(LevelToBitmapDecoder& decoder, int32_t batch_size) { + using BitmapSpanMut = LevelToBitmapDecoder::BitmapSpanMut; + + ReserveValues(batch_size); + const auto write_pos = BitmapSpanMut(data() + values_count() / 8, + static_cast(values_count() % 8)); + const auto decoded = decoder.Decode(batch_size, write_pos); + const auto non_null = + ::arrow::internal::CountSetBits(write_pos.data(), write_pos.bit_start(), decoded); + + ARROW_DCHECK_LE(non_null, decoded); + set_values_count(values_count() + decoded); + return { + .values_read = decoded, + .null_count = static_cast(decoded - non_null), + }; + } + std::shared_ptr ReleaseValues(MemoryPool* pool) { // TODO should we set values_written to zero? auto result = values_; @@ -1134,8 +1154,7 @@ class ValiditySinkBuffer : private ValueSinkCursor { /// If the data page includes repetition and definition levels, we initialize the level /// decoders and return the number of encoded level bytes. /// The return value helps determine the number of bytes in the encoded data. -template -int64_t InitializeV1Levels(const DataPageV1& page, LvlDec& def_dec, LvlDec& rep_dec) { +int64_t InitializeV1Levels(const DataPageV1& page, auto& def_dec, auto& rep_dec) { const auto num_values = static_cast(page.num_values()); const uint8_t* buffer = page.data(); @@ -1165,8 +1184,7 @@ int64_t InitializeV1Levels(const DataPageV1& page, LvlDec& def_dec, LvlDec& rep_ /// If the data page includes repetition and definition levels, we initialize the level /// decoders and return the number of encoded level bytes. /// The return value helps determine the number of bytes in the encoded data. -template -int64_t InitializeV2Levels(const DataPageV2& page, LvlDec& def_dec, LvlDec& rep_dec) { +int64_t InitializeV2Levels(const DataPageV2& page, auto& def_dec, auto& rep_dec) { const auto num_values = static_cast(page.num_values()); const int64_t total_levels_length = @@ -1196,25 +1214,22 @@ int64_t InitializeV2Levels(const DataPageV2& page, LvlDec& def_dec, LvlDec& rep_ } /// Read through the multiple pages of a column chunk. -/// -/// `Traits` must provide two member types: -/// - `DType`: the physical Parquet type of the column; -/// - `level_decoder`: the decoder type used for definition/repetition levels. -/// Each concrete reader defines its own trait struct. template class ColumnChunkReader { public: using DType = typename Traits::DType; - using LvlDec = typename Traits::level_decoder; + using DefLevelDecoder = typename Traits::DefLevelDecoder; + using RepLevelDecoder = typename Traits::RepLevelDecoder; using value_type = typename DType::c_type; ColumnChunkReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, - LvlDec def_levels_decoder, LvlDec rep_levels_decoder) - : descr_(descr), - pool_(pool), - current_decoder_(pool), + DefLevelDecoder def_levels_decoder, + RepLevelDecoder rep_levels_decoder) + : current_decoder_(pool), def_levels_decoder_(std::move(def_levels_decoder)), - rep_levels_decoder_(std::move(rep_levels_decoder)) {} + rep_levels_decoder_(std::move(rep_levels_decoder)), + descr_(descr), + pool_(pool) {} void SetPageReader(std::unique_ptr reader) { current_page_ = nullptr; @@ -1265,15 +1280,15 @@ class ColumnChunkReader { int64_t Skip(int64_t num_values_to_skip); protected: + SkippableTypedDecoder current_decoder_; + DefLevelDecoder def_levels_decoder_; + RepLevelDecoder rep_levels_decoder_; const ColumnDescriptor* descr_; ::arrow::MemoryPool* pool_; - SkippableTypedDecoder current_decoder_; private: using DecoderType = TypedDecoder; - LvlDec def_levels_decoder_; - LvlDec rep_levels_decoder_; // Map of encoding type to the respective decoder object. For example, a // column chunk's data pages may include both dictionary-encoded and // plain-encoded data. @@ -1526,7 +1541,8 @@ int64_t ColumnChunkReader::Skip(int64_t num_values_to_skip) { template struct TypedColumnReaderImplTraits { using DType = D; - using level_decoder = LevelDecoder; + using DefLevelDecoder = LevelDecoder; + using RepLevelDecoder = LevelDecoder; }; template @@ -1763,12 +1779,6 @@ namespace internal { namespace { -/// Whether values of type T can be printed to `std::cout`. -template -concept can_cout = requires(std::ostream& os, const T& value) { - os << value; -}; // NOLINT(readability/braces) - /*********************** * TypedRecordReader * ***********************/ @@ -1776,7 +1786,8 @@ concept can_cout = requires(std::ostream& os, const T& value) { template struct TypedRecordReaderTraits { using DType = D; - using level_decoder = LevelDecoder; + using DefLevelDecoder = LevelDecoder; + using RepLevelDecoder = LevelDecoder; }; template @@ -2576,7 +2587,8 @@ void TypedRecordReader::ResetValues() { template struct RequiredTypedRecordReaderTraits { using DType = D; - using level_decoder = LevelDecoder; + using DefLevelDecoder = LevelDecoder; + using RepLevelDecoder = LevelDecoder; }; // TODO can we reduce some code share with TypedRecordREader ? @@ -2729,6 +2741,226 @@ void RequiredTypedRecordReader::DebugPrintState() { std::cout << std::endl; } +/*********************************** + * FlatOptionalTypedRecordReader * + ***********************************/ + +template +struct FlatOptionalTypedRecordReaderTraits { + using DType = DT; + using DefLevelDecoder = LevelToBitmapDecoder; + using RepLevelDecoder = LevelDecoder; +}; + +template +class FlatOptionalTypedRecordReader + : public ColumnChunkReader>, + virtual public RecordReader { + public: + using T = typename DType::c_type; + using Base = ColumnChunkReader>; + using ValueSink = ValueSinkBuffer; + + FlatOptionalTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool, + bool read_dense_for_nullable, ValueSink value_sink) + : Base(descr, pool, LevelToBitmapDecoder(), LevelDecoder(0)), + value_sink_(std::move(value_sink)), + read_dense_for_nullable_(read_dense_for_nullable) { + ARROW_DCHECK_EQ(descr->max_definition_level(), 1); + ARROW_DCHECK_EQ(descr->max_repetition_level(), 0); + Reset(); + } + + uint8_t* values() const final { return reinterpret_cast(value_sink_.data()); } + + int64_t values_written() const final { return value_sink_.values_count(); } + + int16_t* def_levels() const final { return nullptr; } + + int16_t* rep_levels() const final { return nullptr; } + + int64_t levels_position() const final { return 0; } + + int64_t levels_written() const final { return 0; } + + int64_t null_count() const final { return null_count_; } + + bool nullable_values() const final { return true; } + + bool read_dictionary() const final { return false; } + + bool read_dense_for_nullable() const final { return read_dense_for_nullable_; } + + const void* ReadDictionary(int32_t* dictionary_length) final; + + int64_t ReadRecords(int64_t num_records) final; + + int64_t SkipRecords(int64_t num_records) final { return this->Skip(num_records); } + + std::shared_ptr ReleaseValues() final { + return value_sink_.ReleaseValues(this->pool_); + } + + std::shared_ptr ReleaseIsValid() final; + + void Reserve(int64_t extra_values) final { + ReserveValues(extra_values); + ReserveIsValid(extra_values); + } + + void Reset() final; + + void SetPageReader(std::unique_ptr reader) final { + Base::SetPageReader(std::move(reader)); + // At most one dictionary in Parquet column chunk and it has to be the first page. + if (this->HasPageReader() && this->EnsureDataPage()) { + if (auto* dict = dynamic_cast*>(this->current_decoder_.get())) { + value_sink_.OnNewDictionary(*dict); + } + } + } + + bool HasMoreData() const override { + return Base::HasPageReader(); // Surprising legacy behaviour + } + + const ColumnDescriptor* descr() const final { return this->descr_; } + + void OnNewDictionary(DictDecoder& decoder) { + value_sink_.OnNewDictionary(decoder); + } + + void DebugPrintState() final; + + protected: + auto value_sink() -> ValueSink& { return value_sink_; } + auto value_sink() const -> const ValueSink& { return value_sink_; } + + private: + ValueSink value_sink_; + ValiditySinkBuffer valid_bits_; + int64_t null_count_ = 0; + /// If true, we will not leave any space for the null values in the values. + bool read_dense_for_nullable_ = false; + + void ReserveIsValid(int64_t extra_values); + + void ReserveValues(int64_t extra_values) { value_sink_.ReserveValues(extra_values); } + + void ReadValuesDense(int32_t values_to_read); + + void ReadValuesSpaced(int64_t valid_bits_offset, int32_t values_with_nulls, + int32_t null_count); +}; + +/************************************************** + * FlatOptionalTypedRecordReader Implementation * + **************************************************/ + +template +const void* FlatOptionalTypedRecordReader
::ReadDictionary( + int32_t* dictionary_length) { + if (!this->current_decoder_ && !this->EnsureDataPage()) { + *dictionary_length = 0; + return nullptr; + } + // Verify the current data page is dictionary encoded. + this->CheckEncodingIs(Encoding::RLE_DICTIONARY); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); + const T* dictionary = nullptr; + decoder->GetDictionary(&dictionary, dictionary_length); + return reinterpret_cast(dictionary); +} + +template +int64_t FlatOptionalTypedRecordReader
::ReadRecords(int64_t num_records) { + if (num_records <= 0) { + return 0; + } + + int64_t records_read = 0; + + do { + // Is there more data to read in this row group? + if (!this->EnsureDataPage()) { + break; + } + + const int32_t batch_size = + narrow_min(num_records - records_read, this->available_values_current_page()); + if (read_dense_for_nullable()) { + const auto result = this->def_levels_decoder_.CountUpTo(true, batch_size); + ReadValuesDense(result.matching_count); + records_read += result.processed_count; + this->MarkValuesAsConsumed(result.processed_count); + } else { + const auto old_valid_bits_start = valid_bits_.values_count(); + const auto result = + valid_bits_.ReadFromDecoder(this->def_levels_decoder_, batch_size); + ReadValuesSpaced(old_valid_bits_start, result.values_read, result.null_count); + this->MarkValuesAsConsumed(result.values_read); + records_read += result.values_read; + null_count_ += result.null_count; + } + } while (records_read < num_records); + + return records_read; +} + +template +void FlatOptionalTypedRecordReader
::ReadValuesSpaced(int64_t valid_bits_offset, + int32_t values_with_nulls, + int32_t null_count) { + const auto decoded = value_sink_.ReadValuesSpaced( + *this->current_decoder_.get(), values_with_nulls, null_count, + /* valid_bits= */ valid_bits_.data(), + /* valid_bits_offset= */ valid_bits_offset); + CheckNumberDecoded(decoded, values_with_nulls); +} + +template +void FlatOptionalTypedRecordReader
::ReadValuesDense(int32_t batch_size) { + const auto decoded = + value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size); + CheckNumberDecoded(decoded, batch_size); +} + +template +std::shared_ptr FlatOptionalTypedRecordReader
::ReleaseIsValid() { + if (!read_dense_for_nullable()) { + return valid_bits_.ReleaseValues(this->pool_); + } else { + return nullptr; + } +} + +template +void FlatOptionalTypedRecordReader
::ReserveIsValid(int64_t extra_values) { + if (!read_dense_for_nullable()) { + valid_bits_.ReserveValues(extra_values); + } +} + +template +void FlatOptionalTypedRecordReader
::Reset() { + null_count_ = 0; + // TODO not quite what we want + if (!read_dense_for_nullable()) { + valid_bits_ = ValiditySinkBuffer(this->pool_); + } + + if (values_written() > 0) { + value_sink_.ResetValues(); + } +} + +template +void FlatOptionalTypedRecordReader
::DebugPrintState() { + std::cout << "values: "; + value_sink_.DebugPrintState(); + std::cout << std::endl; +} + /********************* * ArrayValuesSink * *********************/ @@ -3149,23 +3381,25 @@ std::shared_ptr DispatchTypedRecordReader(const ColumnDescriptor* LevelInfo leaf_info, MemoryPool* pool, bool read_dense_for_nullable) { - if (descr->max_definition_level() == 0 && descr->max_repetition_level() == 0) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + if (descr->max_definition_level() == 0 && descr->max_repetition_level() == 0) { using FLBAReader = FLBARecordReader; return std::make_shared(descr, pool); - } else { - using c_type = typename DType::c_type; - using ValueSink = ValueSinkBuffer; - using Reader = RequiredTypedRecordReader; - return std::make_shared(descr, pool, ValueSink(pool)); } - } - if constexpr (std::is_same_v) { using FLBAReader = FLBARecordReader; return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable); } else { using c_type = typename DType::c_type; using ValueSink = ValueSinkBuffer; + if (descr->max_definition_level() == 0 && descr->max_repetition_level() == 0) { + using Reader = RequiredTypedRecordReader; + return std::make_shared(descr, pool, ValueSink(pool)); + } else if (descr->max_definition_level() == 1 && descr->max_repetition_level() == 0 && + descr->schema_node()->is_optional()) { + using Reader = FlatOptionalTypedRecordReader; + return std::make_shared(descr, pool, read_dense_for_nullable, + ValueSink(pool)); + } using Reader = TypedRecordReader; return std::make_shared(descr, leaf_info, pool, read_dense_for_nullable, ValueSink(pool));