Function MakeIntegerBatchGen

cpp/src/arrow/acero/test_util_internal.cc:218–253

}  // namespace

AsyncGenerator<std::optional<ExecBatch>> MakeIntegerBatchGen(
    const std::vector<std::function<int64_t(int)>>& gens,
    const std::shared_ptr<Schema>& schema, int num_batches, int batch_size) {
  struct IntegerBatchGenState {
    IntegerBatchGenState(const std::vector<std::function<int64_t(int)>>& gens,
                         const std::shared_ptr<Schema>& schema, int num_batches,
                         int batch_size)
        : gens(gens), schema(schema), num_batches(num_batches), batch_size(batch_size) {}

    std::optional<ExecBatch> Next() {
      if (batch_index >= num_batches) {
        return std::nullopt;
      }
      Result<ExecBatch> batch_res = MakeIntegerBatch(gens, schema, batch_row, batch_size);
      if (!batch_res.ok()) {
        return std::nullopt;
      }
      ++batch_index;
      batch_row += batch_size;
      return batch_res.ValueOrDie();
    }

    std::vector<std::function<int64_t(int)>> gens;
    std::shared_ptr<Schema> schema;
    int num_batches;
    int batch_size;
    int batch_index = 0;
    int batch_row = 0;
  };
  auto state =
      std::make_shared<IntegerBatchGenState>(gens, schema, num_batches, batch_size);
  return [state]() {
    return DeferNotOk(::arrow::io::default_io_context().executor()->Submit(
        [state]() { return state->Next(); }));
  };
}

BatchesWithSchema MakeBasicBatches() {
  BatchesWithSchema out;