Class RowGroupGenerator

cpp/src/parquet/arrow/reader.cc:1098–1211

Given a file reader and a list of row groups, this is a generator of record batch generators (where each sub-generator is the contents of a single row group).

/// Given a file reader and a list of row groups, this is a generator of record
/// batch generators (where each sub-generator is the contents of a single row group).
class RowGroupGenerator {
 public:
  using RecordBatchGenerator =
      ::arrow::AsyncGenerator<std::shared_ptr<::arrow::RecordBatch>>;

  struct ReadRequest {
    ::arrow::Future<RecordBatchGenerator> read;
    int64_t num_rows;
  };

  explicit RowGroupGenerator(std::shared_ptr<FileReaderImpl> arrow_reader,
                             ::arrow::internal::Executor* cpu_executor,
                             std::vector<int> row_groups, std::vector<int> column_indices,
                             int64_t min_rows_in_flight)
      : arrow_reader_(std::move(arrow_reader)),
        cpu_executor_(cpu_executor),
        row_groups_(std::move(row_groups)),
        column_indices_(std::move(column_indices)),
        min_rows_in_flight_(min_rows_in_flight),
        rows_in_flight_(0),
        index_(0),
        readahead_index_(0) {}

  ::arrow::Future<RecordBatchGenerator> operator()() {
    if (index_ >= row_groups_.size()) {
      return ::arrow::AsyncGeneratorEnd<RecordBatchGenerator>();
    }
    index_++;
    FillReadahead();
    ReadRequest next = std::move(in_flight_reads_.front());
    DCHECK(!in_flight_reads_.empty());
    in_flight_reads_.pop();
    rows_in_flight_ -= next.num_rows;
    return next.read;
  }

 private:
  void FillReadahead() {
    if (min_rows_in_flight_ == 0) {
      // No readahead, fetch the batch when it is asked for
      FetchNext();
    } else {
      while (readahead_index_ < row_groups_.size() &&
             rows_in_flight_ < min_rows_in_flight_) {
        FetchNext();
      }
    }
  }

  void FetchNext() {
    size_t row_group_index = readahead_index_++;
    int row_group = row_groups_[row_group_index];
    std::vector<int> column_indices = column_indices_;
    auto reader = arrow_reader_;
    int64_t num_rows =
        reader->parquet_reader()->metadata()->RowGroup(row_group)->num_rows();
    rows_in_flight_ += num_rows;
    ::arrow::Future<RecordBatchGenerator> row_group_read;