AntoinePrv commented on code in PR #50629:
URL: https://github.com/apache/arrow/pull/50629#discussion_r3957699463
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -1316,969 +1796,1388 @@ 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, bool kReadDictionary>
+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;
-
- if (has_values_to_process()) {
- records_read += ReadRecordData(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->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;
- }
+ int16_t* def_levels() const final { return def_levels_.data(); }
- /// 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());
+ int16_t* rep_levels() const final { return rep_levels_.data(); }
- // 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 final { return kReadDictionary; }
- 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;
-
- int64_t levels_remaining = levels_written_ - gap;
-
- 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));
- };
+ uint8_t* values() const final { return
reinterpret_cast<uint8_t*>(value_sink_.data()); }
- left_shift(def_levels_.get());
+ int64_t values_written() const final { return value_sink_.values_count(); }
- 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, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::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, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::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);
+
+ // For values, we do not have them in buffer, so we will read them and
+ // throw them away.
+ SkipValuesInPage(values_to_read);
+
+ // Mark the levels as read in the underlying column reader.
+ this->MarkValuesAsConsumed(skipped_records);
+
+ return skipped_records;
+}
+
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::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);
+ SkipValuesInPage(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->MarkValuesAsConsumed(clamp_to<int32_t>(levels_position_ -
start_levels_position));
+ // Updated levels_position_ and levels_written().
+ def_levels_.Erase(start_levels_position, levels_position_);
+ rep_levels_.Erase(start_levels_position, levels_position_);
+ levels_position_ = start_levels_position;
+ return skipped_records;
+}
+
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::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 (has_buffered_levels()) {
+ // This updates at_record_start_.
+ skipped_records = DelimitAndSkipRecordsInBuffer(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->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
+ ++skipped_records;
+ at_record_start_ = true;
}
+ break;
+ }
- levels_written_ += batch_size;
- int64_t remaining_records = num_records - skipped_records;
- // This updates at_record_start_.
- skipped_records += DelimitAndSkipRecordsInBuffer(remaining_records);
+ // Read some more levels.
+ const int64_t batch_size_64 =
+ std::min<int64_t>(level_batch_size,
this->available_values_current_page());
+ // available_values_current_page fits in int32_t
+ const auto batch_size = static_cast<int32_t>(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;
}
+ def_levels_.Decode(this->def_levels_decoder_, batch_size);
+ rep_levels_.Decode(this->rep_levels_decoder_, batch_size);
+ const int64_t remaining_records = num_records - skipped_records;
+ // This updates at_record_start_.
+ skipped_records += DelimitAndSkipRecordsInBuffer(remaining_records);
+ }
+
+ return skipped_records;
+}
+
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::SkipValuesInPage(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());
+ }
+}
+
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::SkipRecords(int64_t num_records) {
+ if (num_records == 0) return 0;
+
+ // Top level required field. Number of records equals number of levels,
+ // and there is no read-ahead for levels.
+ if (this->max_rep_level() == 0 && this->max_def_level() == 0) {
+ return this->Skip(num_records);
+ } else if (this->max_rep_level() == 0) {
+ // Non-repeated optional field.
+ // First consume whatever is in the buffer.
+ int64_t 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);
return skipped_records;
}
+ return this->SkipRecordsRepeated(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) {
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::DelimitRecords(int64_t num_records,
+ int64_t* values_seen) {
+ if (ARROW_PREDICT_FALSE(num_records == 0 || !has_buffered_levels())) {
+ *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 exhaust the ColumnChunk
+ int64_t level = levels_position_;
+ if (at_record_start_) {
+ if (ARROW_PREDICT_FALSE(rep_levels[levels_position_] != 0)) {
std::stringstream ss;
- ss << "Could not read and throw away " << num_values << " values";
+ 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 (has_buffered_levels()) {
+ // 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.
+ const 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;
+}
- int64_t SkipRecords(int64_t num_records) override {
- if (num_records == 0) return 0;
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::Reserve(int64_t extra_values) {
+ value_sink_.ReserveValues(extra_values);
+ valid_bits_.ReserveValues(extra_values); // potentially no-op if void
+ def_levels_.ReserveValues(extra_values); // potentially no-op if void
+ rep_levels_.ReserveValues(extra_values); // potentially no-op if void
+}
- // 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);
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::Reset() {
+ null_count_ = 0;
+ value_sink_.ResetValues();
+ valid_bits_.ResetValues(); // potentially no-op if void
+ // Must keep eagerly decoded levels
+ def_levels_.Erase(0, levels_position_); // potentially no-op if void
+ rep_levels_.Erase(0, levels_position_); // potentially no-op if void
+ levels_position_ = 0;
+}
+
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS,
kDic>::SetPageReader(std::unique_ptr<PageReader> 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 = this->current_dict_decoder()) {
+ value_sink_.OnNewDictionary(*dict);
}
- return skipped_records;
}
+}
- // 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_values_to_process() const { return levels_position_ <
levels_written_; }
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::ReadRepeatedRecordsInBuffer(
+ 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 (valid_bits_.is_void()) { // not nullable or read_dense_for_nullable
+ // This is only reading in the current page so this fits in an int32.
+ value_sink_.ReadValuesDense(*this->current_decoder_.get(),
+ clamp_to<int32_t>(*values_to_read));
+ // null_count is always 0 for required.
+ ARROW_DCHECK_EQ(*null_count, 0);
+ } else {
+ ReadSpacedForOptionalOrRepeatedInBuffer(start_levels_position,
values_to_read,
+ null_count);
+ }
+ return records_read;
+}
- std::shared_ptr<ResizableBuffer> 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;
- }
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::ReadOptionalRecordsInBuffer(
+ 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
+ const int64_t records_read =
+ std::min<int64_t>(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()) {
+ ReadDenseForOptionalInBuffer(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 {
+ ReadSpacedForOptionalOrRepeatedInBuffer(start_levels_position,
values_to_read,
+ null_count);
}
+ return records_read;
+}
- std::shared_ptr<ResizableBuffer> ReleaseIsValid() override {
- if (nullable_values()) {
- auto result = valid_bits_;
-
PARQUET_THROW_NOT_OK(result->Resize(bit_util::BytesForBits(values_written_),
- /*shrink_to_fit=*/true));
- valid_bits_ = AllocateBuffer(this->pool_);
- return result;
- } else {
- return nullptr;
- }
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::ReadDenseForOptionalInBuffer(
+ 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());
+ // This is only reading in the current page so this fits in an int32.
+ value_sink_.ReadValuesDense(*this->current_decoder_.get(),
+ clamp_to<int32_t>(*values_to_read));
+}
+
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::ReadSpacedForOptionalOrRepeatedInBuffer(
+ 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.
+ const int64_t valid_bits_offset = valid_bits_.values_count();
+ const auto result =
+ valid_bits_.ReadFromDefLevels(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.
+ value_sink_.ReadValuesSpaced(*this->current_decoder_.get(),
+ clamp_to<int32_t>(result.values_read),
+ clamp_to<int32_t>(*null_count),
+ /* valid_bits= */ valid_bits_.data(),
+ /* valid_bits_offset= */ valid_bits_offset);
+}
+
+template <typename DT, typename VS, bool kDic>
+int64_t TypedRecordReader<DT, VS, kDic>::ReadRecordDataInBuffer(int64_t
num_records) {
+ // 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
+ // 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 = ReadRepeatedRecordsInBuffer(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 = ReadOptionalRecordsInBuffer(num_records, &values_to_read,
&null_count);
+ } else {
+ ARROW_DCHECK(!nullable_values());
+ values_to_read = num_records;
+ // This is only reading in the current page so this fits in an int32.
+ value_sink_.ReadValuesDense(*this->current_decoder_.get(),
+ clamp_to<int32_t>(values_to_read));
+ records_read = num_records;
+ // We don't need to update null_count, since it is 0.
+ }
+
+ ARROW_DCHECK_GE(records_read, 0);
+ ARROW_DCHECK_GE(values_to_read, 0);
+ ARROW_DCHECK_GE(null_count, 0);
+
+ // The values have already been accounted for by the value sink.
+ if (read_dense_for_nullable()) {
+ ARROW_DCHECK_EQ(null_count, 0);
+ } else {
+ null_count_ += null_count;
+ }
+ // Total values, including null spaces, if any
+ if (this->max_def_level() > 0) {
+ // Optional, repeated, or some mix thereof
+ // This is only reading in the current page so this fits in an int32.
+ this->MarkValuesAsConsumed(
+ clamp_to<int32_t>(levels_position_ - start_levels_position));
+ } else {
+ // Flat, non-repeated
+ // This is only reading in the current page so this fits in an int32.
+ this->MarkValuesAsConsumed(clamp_to<int32_t>(values_to_read));
}
- // 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) {
- 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;
+ return records_read;
+}
+
+template <typename DT, typename VS, bool kDic>
+void TypedRecordReader<DT, VS, kDic>::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_;
+
+ 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;
+ }
- // 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;
- }
+ 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] << " ";
}
- // 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;
+ std::cout << std::endl;
}
- void Reserve(int64_t capacity) override {
- ReserveLevels(capacity);
- ReserveValues(capacity);
+ std::cout << "values: ";
+ value_sink_.DebugPrintState();
+ std::cout << std::endl;
+}
+
+/*******************************
+ * RequiredTypedRecordReader *
+ *******************************/
+
+template <typename D>
+struct RequiredTypedRecordReaderTraits {
+ using DType = D;
+ using DefLevelDecoder = NewLevelDecoder;
+ using RepLevelDecoder = NewLevelDecoder;
+};
+
+/// A special type of record reader for required data.
+///
+/// Definition and repetition levels are all null in this case and the data
encoded
+/// correspond directly to the
+template <typename DType, typename ValueSink = ValueSinkBuffer<typename
DType::c_type>,
+ bool kReadDictionary = false>
+class RequiredTypedRecordReader
+ : public ColumnChunkReader<RequiredTypedRecordReaderTraits<DType>>,
+ virtual public RecordReader {
+ public:
+ using T = typename DType::c_type;
+ using Base = ColumnChunkReader<RequiredTypedRecordReaderTraits<DType>>;
+
+ RequiredTypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool,
+ ValueSink value_sink)
+ : Base(descr, pool, NewLevelDecoder(descr->max_definition_level()),
+ NewLevelDecoder(descr->max_repetition_level())),
+ value_sink_(std::move(value_sink)) {
+ RequiredTypedRecordReader::Reset();
+ ARROW_DCHECK_EQ(descr->max_definition_level(), 0);
+ ARROW_DCHECK_EQ(descr->max_repetition_level(), 0);
}
- 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);
+ uint8_t* values() const final { return
reinterpret_cast<uint8_t*>(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 0; }
+
+ bool nullable_values() const final { return false; }
+
+ bool read_dictionary() const final { return kReadDictionary; }
+
+ bool read_dense_for_nullable() const final { return false; }
+
+ const void* ReadDictionary(int32_t* dictionary_length) final {
+ return reinterpret_cast<const
void*>(Base::ReadDictionary(dictionary_length));
}
- 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);
- if (new_levels_capacity > levels_capacity_) {
- constexpr auto kItemSize = static_cast<int64_t>(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;
- }
- }
+ int64_t ReadRecords(int64_t num_records) final;
+
+ int64_t SkipRecords(int64_t num_records) final { return
this->Skip(num_records); }
+
+ std::shared_ptr<ResizableBuffer> ReleaseValues() final {
+ return value_sink_.ReleaseValues(this->pool_);
}
- 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;
- }
- 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<size_t>(valid_bytes_new - valid_bytes_old));
+ std::shared_ptr<ResizableBuffer> ReleaseIsValid() final { return nullptr; }
+
+ void Reserve(int64_t extra_values) final { ReserveValues(extra_values); }
+
+ void Reset() final;
+
+ void SetPageReader(std::unique_ptr<PageReader> 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 = this->current_dict_decoder()) {
+ value_sink_.OnNewDictionary(*dict);
}
}
}
- void Reset() override {
- ResetValues();
+ bool HasMoreData() const override {
+ return Base::HasPageReader(); // Surprising legacy behaviour
+ }
- if (levels_written_ > 0) {
- // Throw away levels from 0 to levels_position_.
- ThrowAwayLevels(0);
- }
+ const ColumnDescriptor* descr() const final { return this->descr_; }
+
+ void DebugPrintState() final;
+
+ protected:
+ auto value_sink() -> ValueSink& { return value_sink_; }
+ auto value_sink() const -> const ValueSink& { return value_sink_; }
+
+ private:
+ ValueSink value_sink_;
+
+ void ReserveValues(int64_t extra_values) {
value_sink_.ReserveValues(extra_values); }
+};
- // Call Finish on the binary builders to reset them
+/**********************************************
+ * RequiredTypedRecordReader Implementation *
+ **********************************************/
+
+template <typename DT, typename VS, bool kDic>
+int64_t RequiredTypedRecordReader<DT, VS, kDic>::ReadRecords(int64_t
num_records) {
+ if (num_records <= 0) {
+ return 0;
}
- void SetPageReader(std::unique_ptr<PageReader> reader) override {
- at_record_start_ = true;
- this->pager_ = std::move(reader);
- ResetDecoders();
+ 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());
+ value_sink_.ReadValuesDense(*this->current_decoder_.get(), batch_size);
+ this->MarkValuesAsConsumed(batch_size);
+
+ records_read += batch_size;
+ } while (records_read < num_records);
+
+ return records_read;
+}
+
+template <typename DT, typename VS, bool kDic>
+void RequiredTypedRecordReader<DT, VS, kDic>::Reset() {
+ if (values_written() > 0) {
+ value_sink_.ResetValues();
}
+}
- bool HasMoreData() const override { return this->pager_ != nullptr; }
+template <typename DT, typename VS, bool kDic>
+void RequiredTypedRecordReader<DT, VS, kDic>::DebugPrintState() {
+ std::cout << "values: ";
+ value_sink_.DebugPrintState();
+ std::cout << std::endl;
+}
- const ColumnDescriptor* descr() const override { return this->descr_; }
+/***********************************
+ * FlatOptionalTypedRecordReader *
+ ***********************************/
- // 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)
{
- uint8_t* valid_bits = valid_bits_->mutable_data();
- const int64_t valid_bits_offset = values_written_;
-
- int64_t num_decoded = this->current_decoder_->DecodeSpaced(
- ValuesHead<T>(), static_cast<int>(values_with_nulls),
- static_cast<int>(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<T>(),
static_cast<int>(values_to_read));
- CheckNumberDecoded(num_decoded, 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) {
- 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 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<int64_t>(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 <typename DT>
+struct FlatOptionalTypedRecordReaderTraits {
+ using DType = DT;
+ using DefLevelDecoder = LevelToBitmapDecoder;
+ using RepLevelDecoder = NewLevelDecoder;
+};
+
+/// A specialized record reader for flat optional data.
+///
+/// In this special case, the max definition level is 1 and these correspond
to the arrow
+/// array we are building. A special level decoder is used to bypass decoding
completely
+/// and only copy the bitmap into the Arrow buffer.
+template <typename DType>
+class FlatOptionalTypedRecordReader
+ : public ColumnChunkReader<FlatOptionalTypedRecordReaderTraits<DType>>,
+ virtual public RecordReader {
+ public:
+ using T = typename DType::c_type;
+ using Base = ColumnChunkReader<FlatOptionalTypedRecordReaderTraits<DType>>;
+ using ValueSink = ValueSinkBuffer<T>;
+
+ FlatOptionalTypedRecordReader(const ColumnDescriptor* descr, MemoryPool*
pool,
+ bool read_dense_for_nullable, ValueSink
value_sink)
+ : Base(descr, pool, LevelToBitmapDecoder(), NewLevelDecoder(0)),
+ value_sink_(std::move(value_sink)) {
+ ARROW_DCHECK_EQ(descr->max_definition_level(), 1);
+ ARROW_DCHECK_EQ(descr->max_repetition_level(), 0);
+ ARROW_DCHECK(descr->schema_node()->is_optional());
+ if (!read_dense_for_nullable) {
+ valid_bits_ = ValiditySinkBuffer::MakeAllocated(this->pool_);
}
- 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;
+ uint8_t* values() const final { return
reinterpret_cast<uint8_t*>(value_sink_.data()); }
+
+ int64_t values_written() const final { return value_sink_.values_count(); }
+
+ int16_t* def_levels() const final { return nullptr; }
Review Comment:
At first I wanted I lazily compute it, then I realize it would not be
possible because this is likely called repeatedly across row groups (so I would
need to do it eagerly).
Since the existing class can still handle the flat optional case, I am
suggesting we enable it behind a flag (default to false) in
`RecordReader::Make` and take the API evolution separately.
--
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]
