ARROW-17318: [C++][Dataset] Support async streaming interface for get…

…ting fragments in Dataset (#13804)

Add `GetFragmentsAsync()` and `GetFragmentsAsyncImpl()`
functions to the generic `Dataset` interface, which
allows to produce fragments in a streamed fashion.

This is one of the prerequisites for making
`FileSystemDataset` to support lazy fragment
processing, which, in turn, can be used to start
scan operations without waiting for the entire
dataset to be discovered.

To aid the transition process of moving to async
implementation in `Dataset`/`AsyncScanner` code,
a default implementation for `GetFragmentsAsyncImpl()`
is provided (yielding a VectorGenerator over
the fragments vector, which is stored by every
implementation of Dataset interface at the moment).

Tests: unit(release)

Signed-off-by: Pavel Solodovnikov <pavel.al.solodovnikov@gmail.com>

Authored-by: Pavel Solodovnikov <pavel.al.solodovnikov@gmail.com>
Signed-off-by: Antoine Pitrou <antoine@python.org>

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

@@ -29,10 +29,16 @@
 #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"
 
 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

@@ -146,6 +146,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

@@ -25,6 +25,7 @@
 #include <optional>
 #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