lidavidm commented on a change in pull request #9947:
URL: https://github.com/apache/arrow/pull/9947#discussion_r611832309
##########
File path: cpp/src/arrow/dataset/file_base.cc
##########
@@ -424,68 +424,21 @@ Status FileSystemDataset::Write(const
FileSystemDatasetWriteOptions& write_optio
std::shared_ptr<Scanner> scanner) {
RETURN_NOT_OK(ValidateBasenameTemplate(write_options.basename_template));
- auto task_group = scanner->options()->TaskGroup();
-
- // Things we'll un-lazy for the sake of simplicity, with the tradeoff they
represent:
- //
- // - Fragment iteration. Keeping this lazy would allow us to start
partitioning/writing
- // any fragments we have before waiting for discovery to complete. This
isn't
- // currently implemented for FileSystemDataset anyway: ARROW-8613
- //
- // - ScanTask iteration. Keeping this lazy would save some unnecessary
blocking when
- // writing Fragments which produce scan tasks slowly. No Fragments do this.
- //
- // NB: neither of these will have any impact whatsoever on the common case
of writing
- // an in-memory table to disk.
- ARROW_ASSIGN_OR_RAISE(auto fragment_it, scanner->GetFragments());
- ARROW_ASSIGN_OR_RAISE(FragmentVector fragments, fragment_it.ToVector());
- ScanTaskVector scan_tasks;
- std::vector<Future<>> scan_futs;
-
- for (const auto& fragment : fragments) {
- auto options = std::make_shared<ScanOptions>(*scanner->options());
- // Avoid contention with multithreaded readers
- options->use_threads = false;
- ARROW_ASSIGN_OR_RAISE(auto scan_task_it,
- Scanner(fragment, std::move(options)).Scan());
- for (auto maybe_scan_task : scan_task_it) {
- ARROW_ASSIGN_OR_RAISE(auto scan_task, maybe_scan_task);
- scan_tasks.push_back(std::move(scan_task));
- }
- }
-
- // 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.
+ // Store a mapping from partitions (represented 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.
WriteState state(write_options);
- for (const auto& scan_task : scan_tasks) {
- if (scan_task->supports_async()) {
- ARROW_ASSIGN_OR_RAISE(auto batches_gen, scan_task->ExecuteAsync());
- 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)));
- }
+ ARROW_ASSIGN_OR_RAISE(auto batch_it, scanner->ScanBatches());
- return Status::OK();
- });
- }
- }
- RETURN_NOT_OK(task_group->Finish());
- auto scan_futs_all_done = AllComplete(scan_futs);
- RETURN_NOT_OK(scan_futs_all_done.status());
+ RETURN_NOT_OK(batch_it.Visit([&state](TaggedRecordBatch tagged_batch) {
Review comment:
I had taken out the equivalent from #9589/ARROW-11797 but I can restore
it.
--
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]
