Function MakeChunkedBatchGenerator

cpp/src/arrow/dataset/dataset_internal.h:140–157

Given a record batch generator, creates a new generator that slices batches so individual batches have at most batch_size rows. The resulting generator is async-reentrant, but does not forward reentrant pulls, so apply readahead before using this helper.

// resulting generator is async-reentrant, but does not forward
// reentrant pulls, so apply readahead before using this helper.
inline RecordBatchGenerator MakeChunkedBatchGenerator(RecordBatchGenerator gen,
                                                      int64_t batch_size) {
  return MakeFlatMappedGenerator(
      std::move(gen),
      [batch_size](const std::shared_ptr<RecordBatch>& batch)
          -> ::arrow::AsyncGenerator<std::shared_ptr<::arrow::RecordBatch>> {
        const int64_t rows = batch->num_rows();
        if (rows <= batch_size) {
          return ::arrow::MakeVectorGenerator<std::shared_ptr<RecordBatch>>({batch});
        }
        std::vector<std::shared_ptr<RecordBatch>> slices;
        slices.reserve(rows / batch_size + (rows % batch_size != 0));
        for (int64_t i = 0; i < rows; i += batch_size) {
          slices.push_back(batch->Slice(i, batch_size));
        }
        return ::arrow::MakeVectorGenerator(std::move(slices));
      });
}

}  // namespace dataset
}  // namespace arrow