ARROW-6964: [C++][Dataset] Add multithread support to Scanner::ToTable

cpp/src/arrow/dataset/scanner.cc

@@ -19,13 +19,16 @@
 
 #include <algorithm>
 #include <memory>
+#include <mutex>
 
 #include "arrow/dataset/dataset.h"
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/filter.h"
 #include "arrow/dataset/scanner_internal.h"
 #include "arrow/table.h"
 #include "arrow/util/iterator.h"
+#include "arrow/util/task_group.h"
+#include "arrow/util/thread_pool.h"
 
 namespace arrow {
 namespace dataset {
@@ -135,6 +138,11 @@ std::shared_ptr<Schema> SchemaFromColumnNames(
   return std::make_shared<Schema>(columns);
 }
 
+Status ScannerBuilder::UseThreads(bool use_threads) {
+  scan_options_->use_threads = use_threads;
+  return Status::OK();
+}
+
 Status ScannerBuilder::Finish(std::unique_ptr<Scanner>* out) const {
   if (has_projection_ && !project_columns_.empty()) {
     scan_options_->projector = std::make_shared<RecordBatchProjector>(
@@ -151,19 +159,54 @@ Status ScannerBuilder::Finish(std::unique_ptr<Scanner>* out) const {
   return Status::OK();
 }
 
-Status Scanner::ToTable(std::shared_ptr<Table>* out) {
+using arrow::internal::TaskGroup;
+
+std::shared_ptr<TaskGroup> Scanner::TaskGroup() const {
+  return options_->use_threads ? TaskGroup::MakeThreaded(context_->thread_pool)
+                               : TaskGroup::MakeSerial();
+}
+
+struct TableAggregator {
+  void Append(std::shared_ptr<RecordBatch> batch) {
+    std::lock_guard<std::mutex> lock(m);
+    batches.emplace_back(std::move(batch));
+  }
+
+  Status Finish(std::shared_ptr<Table>* out) {
+    return Table::FromRecordBatches(batches, out);
+  }
+
+  std::mutex m;
   std::vector<std::shared_ptr<RecordBatch>> batches;
+};
+
+struct ScanTaskPromise {
+  Status operator()() {
+    for (auto maybe_batch : task->Scan()) {
+      ARROW_ASSIGN_OR_RAISE(auto batch, std::move(maybe_batch));
+      aggregator.Append(std::move(batch));
+    }
+
+    return Status::OK();
+  }
+
+  TableAggregator& aggregator;
+  std::shared_ptr<ScanTask> task;
+};
+
+Status Scanner::ToTable(std::shared_ptr<Table>* out) {
+  auto task_group = TaskGroup();
+
+  TableAggregator aggregator;
+  for (auto maybe_scan_task : Scan()) {
+    ARROW_ASSIGN_OR_RAISE(auto scan_task, std::move(maybe_scan_task));
+    task_group->Append(ScanTaskPromise{aggregator, std::move(scan_task)});
+  }
 
-  auto it_scantasks = Scan();
-  RETURN_NOT_OK(it_scantasks.Visit([&batches](std::unique_ptr<ScanTask> task) -> Status {
-    auto it = task->Scan();
-    return it.Visit([&batches](std::shared_ptr<RecordBatch> batch) {
-      batches.push_back(batch);
-      return Status::OK();
-    });
-  }));
+  // Wait for all tasks to complete, or the first error.
+  RETURN_NOT_OK(task_group->Finish());
 
-  return Table::FromRecordBatches(batches, out);
+  return aggregator.Finish(out);
 }
 
 }  // namespace dataset

cpp/src/arrow/dataset/scanner.h

@@ -27,16 +27,22 @@
 #include "arrow/dataset/type_fwd.h"
 #include "arrow/dataset/visibility.h"
 #include "arrow/memory_pool.h"
+#include "arrow/util/thread_pool.h"
 
 namespace arrow {
 
 class Table;
 
+namespace internal {
+class TaskGroup;
+};
+
 namespace dataset {
 
 /// \brief Shared state for a Scan operation
 struct ARROW_DS_EXPORT ScanContext {
   MemoryPool* pool = arrow::default_memory_pool();
+  internal::ThreadPool* thread_pool = arrow::internal::GetCpuThreadPool();
 };
 
 class RecordBatchProjector;
@@ -47,6 +53,10 @@ class ARROW_DS_EXPORT ScanOptions {
 
   static std::shared_ptr<ScanOptions> Defaults();
 
+  // Indicate if the Scanner should make use of the ThreadPool found in the
+  // ScanContext.
+  bool use_threads = false;
+
   // Filter
   std::shared_ptr<Expression> filter;
   // Evaluator for Filter
@@ -109,18 +119,34 @@ Status ScanTaskIteratorFromRecordBatch(std::vector<std::shared_ptr<RecordBatch>>
 ///          yield scan_task
 class ARROW_DS_EXPORT Scanner {
  public:
+  Scanner(DataSourceVector sources, std::shared_ptr<ScanOptions> options,
+          std::shared_ptr<ScanContext> context)
+      : sources_(std::move(sources)),
+        options_(std::move(options)),
+        context_(std::move(context)) {}
+
+  virtual ~Scanner() = default;
+
   /// \brief The Scan operator returns a stream of ScanTask. The caller is
   /// responsible to dispatch/schedule said tasks. Tasks should be safe to run
   /// in a concurrent fashion and outlive the iterator.
   virtual ScanTaskIterator Scan() = 0;
 
-  virtual ~Scanner() = default;
-
   /// \brief Convert a Scanner into a Table.
   ///
+  /// \param[out] out output parameter
+  ///
   /// Use this convenience utility with care. This will serially materialize the
   /// Scan result in memory before creating the Table.
   Status ToTable(std::shared_ptr<Table>* out);
+
+ protected:
+  /// \brief Return a TaskGroup according to ScanContext thread rules.
+  std::shared_ptr<internal::TaskGroup> TaskGroup() const;
+
+  DataSourceVector sources_;
+  std::shared_ptr<ScanOptions> options_;
+  std::shared_ptr<ScanContext> context_;
 };
 
 /// \brief SimpleScanner is a trivial Scanner implementation that flattens
@@ -141,16 +167,9 @@ class ARROW_DS_EXPORT SimpleScanner : public Scanner {
   SimpleScanner(std::vector<std::shared_ptr<DataSource>> sources,
                 std::shared_ptr<ScanOptions> options,
                 std::shared_ptr<ScanContext> context)
-      : sources_(std::move(sources)),
-        options_(std::move(options)),
-        context_(std::move(context)) {}
+      : Scanner(std::move(sources), std::move(options), std::move(context)) {}
 
   ScanTaskIterator Scan() override;
-
- private:
-  std::vector<std::shared_ptr<DataSource>> sources_;
-  std::shared_ptr<ScanOptions> options_;
-  std::shared_ptr<ScanContext> context_;
 };
 
 /// \brief ScannerBuilder is a factory class to construct a Scanner. It is used
@@ -188,6 +207,10 @@ class ARROW_DS_EXPORT ScannerBuilder {
   Status Filter(std::shared_ptr<Expression> filter);
   Status Filter(const Expression& filter);
 
+  /// \brief Indicate if the Scanner should make use of the available
+  ///        ThreadPool found in ScanContext;
+  Status UseThreads(bool use_threads = true);
+
   /// \brief Return the constructed now-immutable Scanner object
   Status Finish(std::unique_ptr<Scanner>* out) const;
 

cpp/src/arrow/dataset/scanner_test.cc

@@ -125,8 +125,14 @@ TEST_F(TestSimpleScanner, ToTable) {
 
   auto scanner = std::make_shared<SimpleScanner>(sources, options_, ctx_);
   std::shared_ptr<Table> actual;
+
   ASSERT_OK(scanner->ToTable(&actual));
+  AssertTablesEqual(*expected, *actual);
 
+  // There is no guarantee on the ordering when using multiple threads, but
+  // since the RecordBatch is always the same it will pass.
+  options_->use_threads = true;
+  ASSERT_OK(scanner->ToTable(&actual));
   AssertTablesEqual(*expected, *actual);
 }