ARROW-17318: [C++][Dataset] Support async streaming interface for getting fragments in Dataset

cpp/src/arrow/dataset/dataset.cc

@@ -15,18 +15,19 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "arrow/dataset/dataset.h"
-
 #include <memory>
 #include <utility>
 
+#include "arrow/dataset/dataset.h"
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/scanner.h"
 #include "arrow/table.h"
+#include "arrow/util/async_generator.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/make_unique.h"
+#include "arrow/util/thread_pool.h"
 
 namespace arrow {
 
@@ -160,6 +161,33 @@ Result<FragmentIterator> Dataset::GetFragments(compute::Expression predicate) {
                                    : MakeEmptyIterator<std::shared_ptr<Fragment>>();
 }
 
+Result<FragmentGenerator> Dataset::GetFragmentsAsync() {
+  return GetFragmentsAsync(compute::literal(true));
+}
+
+Result<FragmentGenerator> Dataset::GetFragmentsAsync(compute::Expression predicate) {
+  ARROW_ASSIGN_OR_RAISE(
+      predicate, SimplifyWithGuarantee(std::move(predicate), partition_expression_));
+  return predicate.IsSatisfiable()
+             ? GetFragmentsAsyncImpl(std::move(predicate),
+                                     arrow::internal::GetCpuThreadPool())
+             : MakeEmptyGenerator<std::shared_ptr<Fragment>>();
+}
+
+// Default impl delegating the work to `GetFragmentsImpl` and wrapping it into
+// BackgroundGenerator/TransferredGenerator, which offloads potentially
+// IO-intensive work to the default IO thread pool and then transfers the control
+// back to the specified executor.
+Result<FragmentGenerator> Dataset::GetFragmentsAsyncImpl(
+    compute::Expression predicate, arrow::internal::Executor* executor) {
+  ARROW_ASSIGN_OR_RAISE(auto iter, GetFragmentsImpl(std::move(predicate)));
+  ARROW_ASSIGN_OR_RAISE(
+      auto background_gen,
+      MakeBackgroundGenerator(std::move(iter), io::default_io_context().executor()));
+  auto transferred_gen = MakeTransferredGenerator(std::move(background_gen), executor);
+  return transferred_gen;
+}
+
 struct VectorRecordBatchGenerator : InMemoryDataset::RecordBatchGenerator {
   explicit VectorRecordBatchGenerator(RecordBatchVector batches)
       : batches_(std::move(batches)) {}

cpp/src/arrow/dataset/dataset.h

@@ -28,11 +28,17 @@
 #include "arrow/compute/exec/expression.h"
 #include "arrow/dataset/type_fwd.h"
 #include "arrow/dataset/visibility.h"
+#include "arrow/util/async_generator_fwd.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/mutex.h"
 #include "arrow/util/optional.h"
 
 namespace arrow {
+
+namespace internal {
+class Executor;
+}  // namespace internal
+
 namespace dataset {
 
 using RecordBatchGenerator = std::function<Future<std::shared_ptr<RecordBatch>>()>;
@@ -134,6 +140,8 @@ class ARROW_DS_EXPORT InMemoryFragment : public Fragment {
 
 /// @}
 
+using FragmentGenerator = AsyncGenerator<std::shared_ptr<Fragment>>;
+
 /// \brief A container of zero or more Fragments.
 ///
 /// A Dataset acts as a union of Fragments, e.g. files deeply nested in a
@@ -148,6 +156,10 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
   Result<FragmentIterator> GetFragments(compute::Expression predicate);
   Result<FragmentIterator> GetFragments();
 
+  /// \brief Async versions of `GetFragments`.
+  Result<FragmentGenerator> GetFragmentsAsync(compute::Expression predicate);
+  Result<FragmentGenerator> GetFragmentsAsync();
+
   const std::shared_ptr<Schema>& schema() const { return schema_; }
 
   /// \brief An expression which evaluates to true for all data viewed by this Dataset.
@@ -174,6 +186,18 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
   Dataset(std::shared_ptr<Schema> schema, compute::Expression partition_expression);
 
   virtual Result<FragmentIterator> GetFragmentsImpl(compute::Expression predicate) = 0;
+  /// \brief Default non-virtual implementation method for the base
+  /// `GetFragmentsAsyncImpl` method, which creates a fragment generator for
+  /// the dataset, possibly filtering results with a predicate (forwarding to
+  /// the synchronous `GetFragmentsImpl` method and moving the computations
+  /// to the background, using the IO thread pool).
+  ///
+  /// Currently, `executor` is always the same as `internal::GetCPUThreadPool()`,
+  /// which means the results from the underlying fragment generator will be
+  /// transfered to the default CPU thread pool. The generator itself is
+  /// offloaded to run on the default IO thread pool.
+  virtual Result<FragmentGenerator> GetFragmentsAsyncImpl(
+      compute::Expression predicate, arrow::internal::Executor* executor);
 
   std::shared_ptr<Schema> schema_;
   compute::Expression partition_expression_ = compute::literal(true);

cpp/src/arrow/dataset/dataset_test.cc

@@ -145,6 +145,34 @@ TEST_F(TestInMemoryDataset, HandlesDifferingSchemas) {
       scanner->ToTable());
 }
 
+TEST_F(TestInMemoryDataset, GetFragmentsSync) {
+  constexpr int64_t kBatchSize = 1024;
+  constexpr int64_t kNumberBatches = 16;
+
+  SetSchema({field("i32", int32()), field("f64", float64())});
+  auto batch = ConstantArrayGenerator::Zeroes(kBatchSize, schema_);
+  auto reader = ConstantArrayGenerator::Repeat(kNumberBatches, batch);
+
+  auto dataset = std::make_shared<InMemoryDataset>(
+      schema_, RecordBatchVector{static_cast<size_t>(kNumberBatches), batch});
+
+  AssertDatasetFragmentsEqual(reader.get(), dataset.get());
+}
+
+TEST_F(TestInMemoryDataset, GetFragmentsAsync) {
+  constexpr int64_t kBatchSize = 1024;
+  constexpr int64_t kNumberBatches = 16;
+
+  SetSchema({field("i32", int32()), field("f64", float64())});
+  auto batch = ConstantArrayGenerator::Zeroes(kBatchSize, schema_);
+  auto reader = ConstantArrayGenerator::Repeat(kNumberBatches, batch);
+
+  auto dataset = std::make_shared<InMemoryDataset>(
+      schema_, RecordBatchVector{static_cast<size_t>(kNumberBatches), batch});
+
+  AssertDatasetAsyncFragmentsEqual(reader.get(), dataset.get());
+}
+
 class TestUnionDataset : public DatasetFixtureMixin {};
 
 TEST_F(TestUnionDataset, ReplaceSchema) {

cpp/src/arrow/dataset/test_util.h

@@ -167,7 +167,7 @@ class DatasetFixtureMixin : public ::testing::Test {
   void AssertFragmentEquals(RecordBatchReader* expected, Fragment* fragment,
                             bool ensure_drained = true) {
     ASSERT_OK_AND_ASSIGN(auto batch_gen, fragment->ScanBatchesAsync(options_));
-    AssertScanTaskEquals(expected, batch_gen);
+    AssertScanTaskEquals(expected, batch_gen, ensure_drained);
 
     if (ensure_drained) {
       EnsureRecordBatchReaderDrained(expected);
@@ -191,6 +191,22 @@ class DatasetFixtureMixin : public ::testing::Test {
     }
   }
 
+  void AssertDatasetAsyncFragmentsEqual(RecordBatchReader* expected, Dataset* dataset,
+                                        bool ensure_drained = true) {
+    ASSERT_OK_AND_ASSIGN(auto predicate, options_->filter.Bind(*dataset->schema()));
+    ASSERT_OK_AND_ASSIGN(auto gen, dataset->GetFragmentsAsync(predicate))
+
+    ASSERT_FINISHES_OK(VisitAsyncGenerator(
+        std::move(gen), [this, expected](const std::shared_ptr<Fragment>& f) {
+          AssertFragmentEquals(expected, f.get(), false /*ensure_drained*/);
+          return Status::OK();
+        }));
+
+    if (ensure_drained) {
+      EnsureRecordBatchReaderDrained(expected);
+    }
+  }
+
   /// \brief Ensure that record batches found in reader are equals to the
   /// record batches yielded by a scanner.
   void AssertScannerEquals(RecordBatchReader* expected, Scanner* scanner,

cpp/src/arrow/util/async_generator.h

@@ -24,6 +24,7 @@
 #include <limits>
 #include <queue>
 
+#include "arrow/util/async_generator_fwd.h"
 #include "arrow/util/async_util.h"
 #include "arrow/util/functional.h"
 #include "arrow/util/future.h"
@@ -66,9 +67,6 @@ namespace arrow {
 //   until all outstanding futures have completed.  Generators that spawn multiple
 //   concurrent futures may need to hold onto an error while other concurrent futures wrap
 //   up.
-template <typename T>
-using AsyncGenerator = std::function<Future<T>()>;
-
 template <typename T>
 struct IterationTraits<AsyncGenerator<T>> {
   /// \brief by default when iterating through a sequence of AsyncGenerator<T>,

cpp/src/arrow/util/async_generator_fwd.h

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <functional>
+
+#include "arrow/type_fwd.h"
+
+namespace arrow {
+
+template <typename T>
+using AsyncGenerator = std::function<Future<T>()>;
+
+template <typename T, typename V>
+class MappingGenerator;
+
+template <typename T, typename ComesAfter, typename IsNext>
+class SequencingGenerator;
+
+template <typename T, typename V>
+class TransformingGenerator;
+
+template <typename T>
+class SerialReadaheadGenerator;
+
+template <typename T>
+class ReadaheadGenerator;
+
+template <typename T>
+class PushGenerator;
+
+template <typename T>
+class MergedGenerator;
+
+template <typename T>
+struct Enumerated;
+
+template <typename T>
+class EnumeratingGenerator;
+
+template <typename T>
+class TransferringGenerator;
+
+template <typename T>
+class BackgroundGenerator;
+
+template <typename T>
+class GeneratorIterator;
+
+template <typename T>
+struct CancellableGenerator;
+
+template <typename T>
+class DefaultIfEmptyGenerator;
+
+}  // namespace arrow