| 2494 | }; |
| 2495 | |
| 2496 | Status ProduceCallback(size_t thread_index) { |
| 2497 | if (done_.load()) return Status::OK(); |
| 2498 | ThreadLocalData& tld = thread_local_data_[thread_index]; |
| 2499 | tld.suppkey_start = rows_generated_.fetch_add(batch_size_); |
| 2500 | if (tld.suppkey_start >= rows_to_generate_) return Status::OK(); |
| 2501 | |
| 2502 | tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start); |
| 2503 | |
| 2504 | tld.batch.resize(SUPPLIER::kNumCols); |
| 2505 | std::fill(tld.batch.begin(), tld.batch.end(), Datum()); |
| 2506 | for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index)); |
| 2507 | |
| 2508 | std::vector<Datum> result(gen_list_.size()); |
| 2509 | for (size_t i = 0; i < gen_list_.size(); i++) { |
| 2510 | int col_idx = gen_list_[i]; |
| 2511 | result[i] = tld.batch[col_idx]; |
| 2512 | } |
| 2513 | ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result))); |
| 2514 | int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_; |
| 2515 | int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1); |
| 2516 | bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1); |
| 2517 | ARROW_RETURN_NOT_OK(output_callback_(std::move(eb))); |
| 2518 | if (is_last_batch) { |
| 2519 | bool expected = false; |
| 2520 | if (done_.compare_exchange_strong(expected, true)) |
| 2521 | ARROW_RETURN_NOT_OK(finished_callback_(batches_outputted_.load())); |
| 2522 | return Status::OK(); |
| 2523 | } |
| 2524 | return schedule_callback_( |
| 2525 | [this](size_t thread_index) { return this->ProduceCallback(thread_index); }); |
| 2526 | } |
| 2527 | |
| 2528 | Status AllocateColumn(size_t thread_index, int column) { |
| 2529 | ThreadLocalData& tld = thread_local_data_[thread_index]; |