westonpace commented on a change in pull request #9892:
URL: https://github.com/apache/arrow/pull/9892#discussion_r610361801
##########
File path: cpp/src/arrow/dataset/file_base.cc
##########
@@ -418,14 +418,46 @@ Status WriteNextBatch(WriteState& state, const
std::shared_ptr<ScanTask>& scan_t
return Status::OK();
}
+Future<> WriteInternal(const ScanOptions& scan_options, WriteState& state,
+ ScanTaskVector scan_tasks, internal::Executor*
cpu_executor) {
+ // Store a mapping from partitions (represened by their formatted partition
expressions)
+ // to a WriteQueue which flushes batches into that partition's output file.
In principle
+ // any thread could produce a batch for any partition, so each task
alternates between
+ // pushing batches and flushing them to disk.
+ std::vector<Future<>> scan_futs;
+ auto task_group = scan_options.TaskGroup();
+
+ for (const auto& scan_task : scan_tasks) {
+ if (scan_task->supports_async()) {
+ ARROW_ASSIGN_OR_RAISE(auto batches_gen,
scan_task->ExecuteAsync(cpu_executor));
+ std::function<Status(std::shared_ptr<RecordBatch> batch)> batch_visitor =
+ [&, scan_task](std::shared_ptr<RecordBatch> batch) {
+ return WriteNextBatch(state, scan_task, std::move(batch));
+ };
+ scan_futs.push_back(VisitAsyncGenerator(batches_gen, batch_visitor));
+ } else {
+ task_group->Append([&, scan_task] {
+ ARROW_ASSIGN_OR_RAISE(auto batches, scan_task->Execute());
+
+ for (auto maybe_batch : batches) {
+ ARROW_ASSIGN_OR_RAISE(auto batch, maybe_batch);
+ RETURN_NOT_OK(WriteNextBatch(state, scan_task, std::move(batch)));
+ }
+
+ return Status::OK();
+ });
+ }
+ }
+ RETURN_NOT_OK(task_group->Finish());
Review comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
