Method ReadNext

cpp/src/arrow/table.cc:687–737

void TableBatchReader::set_chunksize(int64_t chunksize) { max_chunksize_ = chunksize; }

Status TableBatchReader::ReadNext(std::shared_ptr<RecordBatch>* out) {
  if (absolute_row_position_ == table_.num_rows()) {
    *out = nullptr;
    return Status::OK();
  }

  // Determine the minimum contiguous slice across all columns
  int64_t chunksize =
      std::min(table_.num_rows() - absolute_row_position_, max_chunksize_);
  std::vector<const Array*> chunks(table_.num_columns());
  for (int i = 0; i < table_.num_columns(); ++i) {
    auto chunk = column_data_[i]->chunk(chunk_numbers_[i]).get();
    int64_t chunk_remaining = chunk->length() - chunk_offsets_[i];

    if (chunk_remaining < chunksize) {
      chunksize = chunk_remaining;
    }

    chunks[i] = chunk;
  }

  // Slice chunks and advance chunk index as appropriate
  std::vector<std::shared_ptr<ArrayData>> batch_data(table_.num_columns());

  for (int i = 0; i < table_.num_columns(); ++i) {
    // Exhausted chunk
    const Array* chunk = chunks[i];
    const int64_t offset = chunk_offsets_[i];
    std::shared_ptr<ArrayData> slice_data;
    if ((chunk->length() - offset) == chunksize) {
      ++chunk_numbers_[i];
      chunk_offsets_[i] = 0;
      if (offset > 0) {
        // Need to slice
        slice_data = chunk->Slice(offset, chunksize)->data();
      } else {
        // No slice
        slice_data = chunk->data();
      }
    } else {
      chunk_offsets_[i] += chunksize;
      slice_data = chunk->Slice(offset, chunksize)->data();
    }
    batch_data[i] = std::move(slice_data);
  }

  absolute_row_position_ += chunksize;
  *out = RecordBatch::Make(table_.schema(), chunksize, std::move(batch_data));

  return Status::OK();
}

}  // namespace arrow