Copilot commented on code in PR #50629:
URL: https://github.com/apache/arrow/pull/50629#discussion_r4218977687
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -1314,1006 +1687,1475 @@ namespace internal {
namespace {
-template <typename DType>
-class TypedRecordReader : public TypedColumnReaderImpl<DType>,
+/***********************
+ * TypedRecordReader *
+ ***********************/
+
+template <typename D>
+struct TypedRecordReaderTraits {
+ using DType = D;
+ using DefLevelDecoder = NewLevelDecoder;
+ using RepLevelDecoder = NewLevelDecoder;
+};
+
+/// General record reader for a given data type.
+///
+/// This historical class can read all repetition, at the cost of increased
complexity.
+/// The main difficulties are that repeated values will span multiple values
(sometimes
+/// across data pages) in Parquet, while nulls are not written.
+/// Decoding levels is therefore critical to reconstruct the delimitation
across records,
+/// but makes optimizing simple cases harder.
+template <typename DType, typename ValueSink>
+class TypedRecordReader : public
ColumnChunkReader<TypedRecordReaderTraits<DType>>,
virtual public RecordReader {
public:
using T = typename DType::c_type;
- using BASE = TypedColumnReaderImpl<DType>;
- 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) {
- 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();
- }
+ using Base = ColumnChunkReader<TypedRecordReaderTraits<DType>>;
- // 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");
+ TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info,
MemoryPool* pool,
+ bool read_dense_for_nullable, ValueSink value_sink)
+ : Base(descr, pool, NewLevelDecoder(descr->max_definition_level()),
+ NewLevelDecoder(descr->max_repetition_level())),
+ value_sink_(std::move(value_sink)),
+ leaf_info_(leaf_info) {
+ if (!read_dense_for_nullable && nullable_values()) {
+ valid_bits_ = ValiditySinkBuffer::MakeAllocated(this->pool_);
}
- return bytes_for_values;
- }
-
- const void* ReadDictionary(int32_t* dictionary_length) override {
- if (!this->current_decoder_ && !this->HasNextInternal()) {
- *dictionary_length = 0;
- return nullptr;
+ if (this->max_def_level() > 0) {
+ def_levels_ = LevelSinkBuffer::MakeAllocated(pool);
}
- // 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());
+ if (this->max_rep_level() > 0) {
+ rep_levels_ = LevelSinkBuffer::MakeAllocated(pool);
}
- auto decoder =
dynamic_cast<DictDecoder<DType>*>(this->current_decoder_.get());
- const T* dictionary = nullptr;
- decoder->GetDictionary(&dictionary, dictionary_length);
- return reinterpret_cast<const void*>(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;
+ int16_t* def_levels() const final { return def_levels_.data(); }
- if (has_values_to_process()) {
- records_read += ReadRecordData(num_records);
- }
+ int16_t* rep_levels() const final { return rep_levels_.data(); }
- int64_t level_batch_size = std::max<int64_t>(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
- int64_t batch_size =
- std::min(level_batch_size, this->available_values_current_page());
-
- // No more data in column
- if (batch_size == 0) {
- break;
- }
+ int64_t levels_position() const final { return levels_position_; }
- if (this->max_def_level() > 0) {
- ReserveLevels(batch_size);
+ int64_t levels_written() const final { return def_levels_.values_count(); }
- int16_t* def_levels = this->def_levels() + levels_written_;
- int16_t* rep_levels = this->rep_levels() + levels_written_;
+ int64_t null_count() const final { return null_count_; }
- 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);
- }
- }
+ bool nullable_values() const final { return leaf_info_.HasNullableValues(); }
- levels_written_ += batch_size;
- 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);
- records_read += ReadRecordData(batch_size);
- }
- }
+ bool read_dictionary() const override { return false; }
- return records_read;
+ bool read_dense_for_nullable() const final {
+ // false for required types regardless of input
+ return nullable_values() && valid_bits_.is_void();
}
- // 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;
+ uint8_t* values() const final { return
reinterpret_cast<uint8_t*>(value_sink_.data()); }
- int64_t levels_remaining = levels_written_ - gap;
+ int64_t values_written() const final { return value_sink_.values_count(); }
- auto left_shift = [&](::arrow::ResizableBuffer* buffer) {
- auto* data = buffer->mutable_data_as<int16_t>();
- 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;
+ const void* ReadDictionary(int32_t* dictionary_length) final {
+ return reinterpret_cast<const
void*>(Base::ReadDictionary(dictionary_length));
}
+ int64_t ReadRecords(int64_t num_records) override;
+
// 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 SkipRecordsRepeated(int64_t num_records);
- int64_t level_batch_size =
- std::max<int64_t>(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;
- }
+ // Skip 'num_values' values from the current page.
+ void SkipValuesInPage(int64_t num_values);
- // Read some more levels.
- int64_t batch_size =
- std::min(level_batch_size, this->available_values_current_page());
- // 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;
+ 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
+ bool has_buffered_levels() const {
+ ARROW_DCHECK_LE(levels_position_, levels_written());
+ return levels_position_ < levels_written();
+ }
+
+ std::shared_ptr<ResizableBuffer> ReleaseValues() override {
+ return value_sink_.ReleaseValues(this->pool_);
+ }
+
+ std::shared_ptr<ResizableBuffer> ReleaseIsValid() final {
+ return valid_bits_.ReleaseValues(this->pool_); // nullptr if void
+ }
+
+ // Process written repetition/definition levels to reach the end of
+ // records. Only used for repeated fields.
+ // Process no more levels than necessary to delimit the indicated
+ // 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);
+
+ void Reserve(int64_t capacity) override;
+
+ void Reset() override;
+
+ void SetPageReader(std::unique_ptr<PageReader> reader) override;
+
+ bool HasMoreData() const override {
+ return Base::HasPageReader(); // Surprising legacy behaviour
+ }
+
+ const ColumnDescriptor* descr() const override { return this->descr_; }
+
+ // Reads repeated records from the buffered levels and returns number of
records
+ // read. Fills in values_to_read and null_count.
+ int64_t ReadRepeatedRecordsInBuffer(int64_t num_records, int64_t*
values_to_read,
+ int64_t* null_count);
+
+ // Reads optional records from the buffered levels and returns number of
records
+ // read. Fills in values_to_read and null_count.
+ int64_t ReadOptionalRecordsInBuffer(int64_t num_records, int64_t*
values_to_read,
+ int64_t* null_count);
+
+ // Reads dense for optional records. First it figures out how many values to
+ // read.
+ void ReadDenseForOptionalInBuffer(int64_t start_levels_position,
+ int64_t* values_to_read);
+
+ // Reads spaced for optional or repeated fields.
+ void ReadSpacedForOptionalOrRepeatedInBuffer(int64_t start_levels_position,
+ int64_t* values_to_read,
+ int64_t* null_count);
+
+ // Read records from the buffered levels.
+ // Return number of logical records read.
+ // Updates levels_position_, values_written_, and null_count_.
+ int64_t ReadRecordDataInBuffer(int64_t num_records);
+
+ void DebugPrintState() override;
+
+ protected:
+ 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.
+ ValiditySinkBuffer valid_bits_ = ValiditySinkBuffer::MakeVoid();
+ LevelInfo leaf_info_;
+
+ /// \brief Buffer for definition levels.
+ ///
+ /// May contain eagerly decoded levels as required to figure out record
boundaries for
+ /// repeated fields. `level_position_` is the number of level processed for
the current
+ /// decoded values. For flat required fields, `def_levels_` and
`rep_levels_` are
+ /// not populated nor allocated.
+ /// `def_levels_` and `rep_levels_` must be of the same size if present.
+ LevelSinkBuffer def_levels_ = LevelSinkBuffer::MakeVoid();
+ /// \brief Buffer for repetition levels. Only populated for repeated fields.
+ LevelSinkBuffer rep_levels_ = LevelSinkBuffer::MakeVoid();
+ /// \brief Position of the next level that should be consumed.
+ int64_t levels_position_ = 0;
+
+ int64_t null_count_ = 0;
+
+ bool at_record_start_ = true;
+};
+
+/**************************************
+ * TypedRecordReader Implementation *
+ **************************************/
+
+template <typename DT, typename VS>
+int64_t TypedRecordReader<DT, VS>::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_buffered_levels()) {
+ records_read += ReadRecordDataInBuffer(num_records);
+ }
+
+ int64_t level_batch_size = std::max<int64_t>(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->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
+ // the row group
+ ++records_read;
+ at_record_start_ = true;
}
+ 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);
+ /// We perform multiple batch reads until we either exhaust the row group
+ /// or observe the desired number of records
+ const int32_t batch_size =
+ narrow_min(level_batch_size, this->available_values_current_page());
- int16_t* def_levels = this->def_levels() + levels_written_;
- int16_t* rep_levels = this->rep_levels() + levels_written_;
+ // No more data in column
+ if (batch_size == 0) {
+ break;
+ }
- if (this->ReadDefinitionLevels(batch_size, def_levels) != batch_size) {
- throw ParquetException(kErrorRepDefLevelNotMatchesNumValues);
+ if (this->max_def_level() > 0) {
+ def_levels_.Decode(this->def_levels_decoder_, batch_size);
+ if (this->max_rep_level() > 0) {
+ rep_levels_.Decode(this->rep_levels_decoder_, batch_size);
}
- if (this->ReadRepetitionLevels(batch_size, rep_levels) != batch_size) {
- throw ParquetException(kErrorRepDefLevelNotMatchesNumValues);
+ records_read += ReadRecordDataInBuffer(num_records - records_read);
+ } else {
+ // No repetition and definition levels, we can read values directly
+ const auto count = narrow_min(num_records - records_read, batch_size);
+ records_read += ReadRecordDataInBuffer(count);
+ }
+ }
+
+ return records_read;
+}
+
+template <typename DT, typename VS>
+int64_t TypedRecordReader<DT, VS>::SkipRecordsInBufferNonRepeated(int64_t
num_records) {
+ ARROW_DCHECK_EQ(this->max_rep_level(), 0);
+ if (!this->has_buffered_levels() || num_records == 0) return 0;
+
+ const int64_t remaining_records_64 = levels_written() - levels_position_;
+ ARROW_DCHECK_LE(remaining_records_64, std::numeric_limits<int32_t>::max());
+ const auto remaining_records = static_cast<int32_t>(remaining_records_64);
+ const int32_t skipped_records = narrow_min(num_records, remaining_records);
+ const int64_t remaining_levels_pos = 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.
+ const auto values_to_read = static_cast<int32_t>(
+ std::count(def_levels() + levels_position_, def_levels() +
remaining_levels_pos,
+ 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.
+ def_levels_.Erase(levels_position_, remaining_levels_pos);
Review Comment:
`Erase` compacts the unread levels down to `levels_position_` and reduces
`levels_written()`, but this method leaves `levels_position_` unchanged. If
records were already consumed before this skip (so the position is nonzero),
the cursor then points past the compacted buffer; subsequent
`has_buffered_levels()`/skips can miss or overrun levels. Reset the cursor to
the old start position after erasing, as `DelimitAndSkipRecordsInBuffer` does
below.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
