ARROW-14192: [C++][Dataset] Backpressure broken on ordered scans

cpp/src/arrow/dataset/scanner.cc

@@ -441,7 +441,8 @@ class AsyncScanner : public Scanner, public std::enable_shared_from_this<AsyncSc
  Result<TaggedRecordBatchGenerator> ScanBatchesAsync(Executor* executor);
  Future<> VisitBatchesAsync(std::function<Status(TaggedRecordBatch)> visitor,
  Executor* executor);
- Result<EnumeratedRecordBatchGenerator> ScanBatchesUnorderedAsync(Executor* executor);
+ Result<EnumeratedRecordBatchGenerator> ScanBatchesUnorderedAsync(
+ Executor* executor, bool sequence_fragments = false);
  Future<std::shared_ptr<Table>> ToTableAsync(Executor* executor);
 
  Result<FragmentGenerator> GetFragments() const;
@@ -587,7 +588,7 @@ Result<EnumeratedRecordBatch> ToEnumeratedRecordBatch(
 }
 
 Result<EnumeratedRecordBatchGenerator> AsyncScanner::ScanBatchesUnorderedAsync(
- Executor* cpu_executor) {
+ Executor* cpu_executor, bool sequence_fragments) {
  if (!scan_options_->use_threads) {
  cpu_executor = nullptr;
  }
@@ -608,7 +609,8 @@ Result<EnumeratedRecordBatchGenerator> AsyncScanner::ScanBatchesUnorderedAsync(
  RETURN_NOT_OK(
  compute::Declaration::Sequence(
  {
- {"scan", ScanNodeOptions{dataset_, scan_options_, backpressure.toggle}},
+ {"scan", ScanNodeOptions{dataset_, scan_options_, backpressure.toggle,
+ sequence_fragments}},
  {"filter", compute::FilterNodeOptions{scan_options_->filter}},
  {"augmented_project",
  compute::ProjectNodeOptions{std::move(exprs), std::move(names)}},
@@ -649,7 +651,8 @@ Result<TaggedRecordBatchGenerator> AsyncScanner::ScanBatchesAsync() {
 
 Result<TaggedRecordBatchGenerator> AsyncScanner::ScanBatchesAsync(
  Executor* cpu_executor) {
- ARROW_ASSIGN_OR_RAISE(auto unordered, ScanBatchesUnorderedAsync(cpu_executor));
+ ARROW_ASSIGN_OR_RAISE(auto unordered, ScanBatchesUnorderedAsync(
+ cpu_executor, /*sequence_fragments=*/true));
  // We need an initial value sentinel, so we use one with fragment.index < 0
  auto is_before_any = [](const EnumeratedRecordBatch& batch) {
  return batch.fragment.index < 0;
@@ -1143,6 +1146,7 @@ Result<compute::ExecNode*> MakeScanNode(compute::ExecPlan* plan,
  auto scan_options = scan_node_options.scan_options;
  auto dataset = scan_node_options.dataset;
  const auto& backpressure_toggle = scan_node_options.backpressure_toggle;
+ bool require_sequenced_output = scan_node_options.require_sequenced_output;
 
  if (!scan_options->use_async) {
  return Status::NotImplemented("ScanNodes without asynchrony");
@@ -1175,8 +1179,15 @@ Result<compute::ExecNode*> MakeScanNode(compute::ExecPlan* plan,
  ARROW_ASSIGN_OR_RAISE(auto batch_gen_gen,
  FragmentsToBatches(std::move(fragment_gen), scan_options));
 
- auto merged_batch_gen =
- MakeMergedGenerator(std::move(batch_gen_gen), scan_options->fragment_readahead);
+ AsyncGenerator<EnumeratedRecordBatch> merged_batch_gen;
+ if (require_sequenced_output) {
+ ARROW_ASSIGN_OR_RAISE(merged_batch_gen,
+ MakeSequencedMergedGenerator(std::move(batch_gen_gen),
+ scan_options->fragment_readahead));
+ } else {
+ merged_batch_gen =
+ MakeMergedGenerator(std::move(batch_gen_gen), scan_options->fragment_readahead);
+ }
 
  auto batch_gen = MakeReadaheadGenerator(std::move(merged_batch_gen),
  scan_options->fragment_readahead);

cpp/src/arrow/dataset/scanner.h

@@ -421,14 +421,17 @@ class ARROW_DS_EXPORT ScanNodeOptions : public compute::ExecNodeOptions {
  public:
  explicit ScanNodeOptions(
  std::shared_ptr<Dataset> dataset, std::shared_ptr<ScanOptions> scan_options,
- std::shared_ptr<util::AsyncToggle> backpressure_toggle = NULLPTR)
+ std::shared_ptr<util::AsyncToggle> backpressure_toggle = NULLPTR,
+ bool require_sequenced_output = false)
  : dataset(std::move(dataset)),
  scan_options(std::move(scan_options)),
- backpressure_toggle(std::move(backpressure_toggle)) {}
+ backpressure_toggle(std::move(backpressure_toggle)),
+ require_sequenced_output(require_sequenced_output) {}
 
  std::shared_ptr<Dataset> dataset;
  std::shared_ptr<ScanOptions> scan_options;
  std::shared_ptr<util::AsyncToggle> backpressure_toggle;
+ bool require_sequenced_output;
 };
 
 /// @}

cpp/src/arrow/dataset/scanner_test.cc

@@ -1027,12 +1027,12 @@ class TestBackpressure : public ::testing::Test {
  return sum;
  }
 
- void Finish(AsyncGenerator<EnumeratedRecordBatch> gen) {
+ template <typename T>
+ void Finish(AsyncGenerator<T> gen) {
  for (const auto& controlled_fragment : controlled_fragments_) {
  controlled_fragment->Finish();
  }
- ASSERT_FINISHES_OK(VisitAsyncGenerator(
- gen, [](EnumeratedRecordBatch batch) { return Status::OK(); }));
+ ASSERT_FINISHES_OK(VisitAsyncGenerator(gen, [](T batch) { return Status::OK(); }));
  }
 
  std::shared_ptr<Schema> schema_ = schema({field("values", int32())});
@@ -1063,6 +1063,24 @@ TEST_F(TestBackpressure, ScanBatchesUnordered) {
  Finish(std::move(gen));
 }
 
+TEST_F(TestBackpressure, ScanBatchesOrdered) {
+ std::shared_ptr<Scanner> scanner = MakeScanner();
+ EXPECT_OK_AND_ASSIGN(AsyncGenerator<TaggedRecordBatch> gen,
+ scanner->ScanBatchesAsync());
+ // This future never actually finishes because we only emit the first batch so far and
+ // the scanner delays by one batch. It is enough to start the system pumping though so
+ // we don't need it to finish.
+ Future<TaggedRecordBatch> fut = gen();
+
+ // See note on other test
+ GetCpuThreadPool()->WaitForIdle();
+ // Worst case we read in the entire set of initial batches
+ ASSERT_LE(TotalBatchesRead(), NBATCHES * (NFRAGMENTS - 1) + 1);
+
+ DeliverAdditionalBatches();
+ Finish(std::move(gen));
+}
+
 struct BatchConsumer {
  explicit BatchConsumer(EnumeratedRecordBatchGenerator generator)
  : generator(std::move(generator)), next() {}

cpp/src/arrow/util/async_generator.h

@@ -111,14 +111,14 @@ Future<> VisitAsyncGenerator(AsyncGenerator<T> generator, Visitor visitor) {
  return Loop(LoopBody{std::move(generator), std::move(visitor)});
 }
 
-/// \brief Waits for an async generator to complete, discarding results.
+/// \brief Wait for an async generator to complete, discarding results.
 template <typename T>
 Future<> DiscardAllFromAsyncGenerator(AsyncGenerator<T> generator) {
  std::function<Status(T)> visitor = [](const T&) { return Status::OK(); };
  return VisitAsyncGenerator(generator, visitor);
 }
 
-/// \brief Collects the results of an async generator into a vector
+/// \brief Collect the results of an async generator into a vector
 template <typename T>
 Future<std::vector<T>> CollectAsyncGenerator(AsyncGenerator<T> generator) {
  auto vec = std::make_shared<std::vector<T>>();
@@ -258,7 +258,7 @@ class MappingGenerator {
  std::shared_ptr<State> state_;
 };
 
-/// \brief Creates a generator that will apply the map function to each element of
+/// \brief Create a generator that will apply the map function to each element of
 /// source. The map function is not called on the end token.
 ///
 /// Note: This function makes a copy of `map` for each item
@@ -278,7 +278,7 @@ AsyncGenerator<V> MakeMappedGenerator(AsyncGenerator<T> source_generator, MapFn
  return MappingGenerator<T, V>(std::move(source_generator), MapCallback{std::move(map)});
 }
 
-/// \brief Creates a generator that will apply the map function to
+/// \brief Create a generator that will apply the map function to
 /// each element of source. The map function is not called on the end
 /// token. The result of the map function should be another
 /// generator; all these generators will then be flattened to produce
@@ -417,7 +417,7 @@ class SequencingGenerator {
  const std::shared_ptr<State> state_;
 };
 
-/// \brief Buffers an AsyncGenerator to return values in sequence order ComesAfter
+/// \brief Buffer an AsyncGenerator to return values in sequence order ComesAfter
 /// and IsNext determine the sequence order.
 ///
 /// ComesAfter should be a BinaryPredicate that only returns true if a comes after b
@@ -532,7 +532,7 @@ class TransformingGenerator {
  std::shared_ptr<TransformingGeneratorState> state_;
 };
 
-/// \brief Transforms an async generator using a transformer function returning a new
+/// \brief Transform an async generator using a transformer function returning a new
 /// AsyncGenerator
 ///
 /// The transform function here behaves exactly the same as the transform function in
@@ -686,7 +686,7 @@ class FutureFirstGenerator {
  std::shared_ptr<State> state_;
 };
 
-/// \brief Transforms a Future<AsyncGenerator<T>> into an AsyncGenerator<T>
+/// \brief Transform a Future<AsyncGenerator<T>> into an AsyncGenerator<T>
 /// that waits for the future to complete as part of the first item.
 ///
 /// This generator is not async-reentrant (even if the generator yielded by future is)
@@ -697,7 +697,7 @@ AsyncGenerator<T> MakeFromFuture(Future<AsyncGenerator<T>> future) {
  return FutureFirstGenerator<T>(std::move(future));
 }
 
-/// \brief Creates a generator that will pull from the source into a queue. Unlike
+/// \brief Create a generator that will pull from the source into a queue. Unlike
 /// MakeReadaheadGenerator this will not pull reentrantly from the source.
 ///
 /// The source generator does not need to be async-reentrant
@@ -711,6 +711,35 @@ AsyncGenerator<T> MakeSerialReadaheadGenerator(AsyncGenerator<T> source_generato
  return SerialReadaheadGenerator<T>(std::move(source_generator), max_readahead);
 }
 
+/// \brief Create a generator that immediately pulls from the source
+///
+/// Typical generators do not pull from their source until they themselves
+/// are pulled. This generator does not follow that convention and will call
+/// generator() once before it returns. The returned generator will otherwise
+/// mirror the source.
+///
+/// This generator forwards aysnc-reentrant pressure to the source
+/// This generator buffers one item (the first result) until it is delivered.
+template <typename T>
+AsyncGenerator<T> MakeAutoStartingGenerator(AsyncGenerator<T> generator) {
+ struct AutostartGenerator {
+ Future<T> operator()() {
+ if (first_future->is_valid()) {
+ Future<T> result = *first_future;
+ *first_future = Future<T>();
+ return result;
+ }
+ return source();
+ }
+
+ std::shared_ptr<Future<T>> first_future;
+ AsyncGenerator<T> source;
+ };
+
+ std::shared_ptr<Future<T>> first_future = std::make_shared<Future<T>>(generator());
+ return AutostartGenerator{std::move(first_future), std::move(generator)};
+}
+
 /// \see MakeReadaheadGenerator
 template <typename T>
 class ReadaheadGenerator {
@@ -919,7 +948,7 @@ class PushGenerator {
  const std::shared_ptr<State> state_;
 };
 
-/// \brief Creates a generator that pulls reentrantly from a source
+/// \brief Create a generator that pulls reentrantly from a source
 /// This generator will pull reentrantly from a source, ensuring that max_readahead
 /// requests are active at any given time.
 ///
@@ -1137,7 +1166,7 @@ class MergedGenerator {
  std::shared_ptr<State> state_;
 };
 
-/// \brief Creates a generator that takes in a stream of generators and pulls from up to
+/// \brief Create a generator that takes in a stream of generators and pulls from up to
 /// max_subscriptions at a time
 ///
 /// Note: This may deliver items out of sequence. For example, items from the third
@@ -1156,7 +1185,24 @@ AsyncGenerator<T> MakeMergedGenerator(AsyncGenerator<AsyncGenerator<T>> source,
  return MergedGenerator<T>(std::move(source), max_subscriptions);
 }
 
-/// \brief Creates a generator that takes in a stream of generators and pulls from each
+template <typename T>
+Result<AsyncGenerator<T>> MakeSequencedMergedGenerator(
+ AsyncGenerator<AsyncGenerator<T>> source, int max_subscriptions) {
+ if (max_subscriptions < 0) {
+ return Status::Invalid("max_subscriptions must be a positive integer");
+ }
+ if (max_subscriptions == 1) {
+ return Status::Invalid("Use MakeConcatenatedGenerator if max_subscriptions is 1");
+ }
+ AsyncGenerator<AsyncGenerator<T>> autostarting_source = MakeMappedGenerator(
+ std::move(source),
+ [](const AsyncGenerator<T>& sub) { return MakeAutoStartingGenerator(sub); });
+ AsyncGenerator<AsyncGenerator<T>> sub_readahead =
+ MakeSerialReadaheadGenerator(std::move(autostarting_source), max_subscriptions - 1);
+ return MakeConcatenatedGenerator(std::move(sub_readahead));
+}
+
+/// \brief Create a generator that takes in a stream of generators and pulls from each
 /// one in sequence.
 ///
 /// This generator is async-reentrant but will never pull from source reentrantly and
@@ -1224,7 +1270,7 @@ class EnumeratingGenerator {
  std::shared_ptr<State> state_;
 };
 
-/// Wraps items from a source generator with positional information
+/// Wrap items from a source generator with positional information
 ///
 /// When used with MakeMergedGenerator and MakeSequencingGenerator this allows items to be
 /// processed in a "first-available" fashion and later resequenced which can reduce the
@@ -1260,7 +1306,7 @@ class TransferringGenerator {
  internal::Executor* executor_;
 };
 
-/// \brief Transfers a future to an underlying executor.
+/// \brief Transfer a future to an underlying executor.
 ///
 /// Continuations run on the returned future will be run on the given executor
 /// if they cannot be run synchronously.
@@ -1506,7 +1552,7 @@ class BackgroundGenerator {
 constexpr int kDefaultBackgroundMaxQ = 32;
 constexpr int kDefaultBackgroundQRestart = 16;
 
-/// \brief Creates an AsyncGenerator<T> by iterating over an Iterator<T> on a background
+/// \brief Create an AsyncGenerator<T> by iterating over an Iterator<T> on a background
 /// thread
 ///
 /// The parameter max_q and q_restart control queue size and background thread task
@@ -1554,14 +1600,14 @@ class GeneratorIterator {
  AsyncGenerator<T> source_;
 };
 
-/// \brief Converts an AsyncGenerator<T> to an Iterator<T> by blocking until each future
+/// \brief Convert an AsyncGenerator<T> to an Iterator<T> which blocks until each future
 /// is finished
 template <typename T>
 Iterator<T> MakeGeneratorIterator(AsyncGenerator<T> source) {
  return Iterator<T>(GeneratorIterator<T>(std::move(source)));
 }
 
-/// \brief Adds readahead to an iterator using a background thread.
+/// \brief Add readahead to an iterator using a background thread.
 ///
 /// Under the hood this is converting the iterator to a generator using
 /// MakeBackgroundGenerator, adding readahead to the converted generator with
@@ -1633,7 +1679,7 @@ AsyncGenerator<T> MakeFailingGenerator(const Result<T>& result) {
  return MakeFailingGenerator<T>(result.status());
 }
 
-/// \brief Prepends initial_values onto a generator
+/// \brief Prepend initial_values onto a generator
 ///
 /// This generator is async-reentrant but will buffer requests and will not
 /// pull from following_values async-reentrantly.
@@ -1659,7 +1705,7 @@ struct CancellableGenerator {
  StopToken stop_token;
 };
 
-/// \brief Allows an async generator to be cancelled
+/// \brief Allow an async generator to be cancelled
 ///
 /// This generator is async-reentrant
 template <typename T>
@@ -1698,7 +1744,7 @@ struct PauseableGenerator {
  std::shared_ptr<PauseableGeneratorState> state_;
 };
 
-/// \brief Allows an async generator to be paused
+/// \brief Allow an async generator to be paused
 ///
 /// This generator is NOT async-reentrant and calling it in an async-reentrant fashion
 /// may lead to items getting reordered (and potentially truncated if the end token is