GH-46421: [C++][Acero] Asofjoin respect PauseProducing from downstream.

cpp/src/arrow/acero/asof_join_node.cc

@@ -1063,15 +1063,23 @@ class AsofJoinNode : public ExecNode {
   }
 
   bool Process() {
-    std::lock_guard<std::mutex> guard(gate_);
-    if (!CheckEnded()) {
-      return false;
-    }
-
     // Process batches while we have data
     for (;;) {
-      Result<std::shared_ptr<RecordBatch>> result = ProcessInner();
+      Future<> to_wait;
+      {
+        std::lock_guard<std::mutex> lg(backpressure_mutex_);
+        to_wait = backpressure_future_;
+      }
+      to_wait.Wait();
 
+      Result<std::shared_ptr<RecordBatch>> result;
+      {
+        std::lock_guard<std::mutex> guard(gate_);
+        if (!CheckEnded()) {
+          return false;
+        }
+        result = ProcessInner();
+      }
       if (result.ok()) {
         auto out_rb = *result;
         if (!out_rb) break;
@@ -1524,8 +1532,51 @@ class AsofJoinNode : public ExecNode {
     return Status::OK();
   }
 
-  void PauseProducing(ExecNode* output, int32_t counter) override {}
-  void ResumeProducing(ExecNode* output, int32_t counter) override {}
+  void PauseProducing(ExecNode* output, int32_t counter) override {
+    std::lock_guard<std::mutex> lg(backpressure_mutex_);
+    if (counter <= last_backpressure_counter_) {
+      return;
+    }
+    last_backpressure_counter_ = counter;
+    if (!backpressure_future_.is_finished()) {
+      // Could happen if we get something like Pause(1) Pause(3) Resume(2)
+      return;
+    }
+    backpressure_future_ = Future<>::Make();
+  }
+  void ResumeProducing(ExecNode* output, int32_t counter) override {
+    Future<> to_finish;
+    {
+      std::lock_guard<std::mutex> lg(backpressure_mutex_);
+      if (counter <= last_backpressure_counter_) {
+        return;
+      }
+      last_backpressure_counter_ = counter;
+      if (backpressure_future_.is_finished()) {
+        return;
+      }
+      to_finish = backpressure_future_;
+      backpressure_future_ = Future<>::MakeFinished();
+    }
+    to_finish.MarkFinished();
+  }
+
+  Status StopProducing() override {
+    // GH-35837: ensure node is not paused
+    Future<> to_finish;
+    {
+      std::lock_guard<std::mutex> lg(backpressure_mutex_);
+      if (!backpressure_future_.is_finished()) {
+        to_finish = backpressure_future_;
+        backpressure_future_ = Future<>::MakeFinished();
+      }
+    }
+    if (to_finish.is_valid()) {
+      to_finish.MarkFinished();
+    }
+    // only then stop
+    return ExecNode::StopProducing();
+  }
 
   Status StopProducingImpl() override {
 #ifdef ARROW_ENABLE_THREADING
@@ -1553,6 +1604,11 @@ class AsofJoinNode : public ExecNode {
   // Each input state corresponds to an input table
   std::vector<std::unique_ptr<InputState>> state_;
   std::mutex gate_;
+
+  std::mutex backpressure_mutex_;
+  std::atomic<int32_t> last_backpressure_counter_{0};
+  Future<> backpressure_future_ = Future<>::MakeFinished();
+
   TolType tolerance_;
 #ifndef NDEBUG
   std::ostream* debug_os_;

cpp/src/arrow/acero/asof_join_node_test.cc

@@ -44,6 +44,7 @@
 #include "arrow/compute/cast.h"
 #include "arrow/compute/row/row_encoder_internal.h"
 #include "arrow/compute/test_util_internal.h"
+#include "arrow/io/util_internal.h"
 #include "arrow/testing/generator.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/matchers.h"
@@ -1412,15 +1413,28 @@ struct BackpressureCountingNode : public MapNode {
   Result<ExecBatch> ProcessBatch(ExecBatch batch) override { return batch; }
 
   void PauseProducing(ExecNode* output, int32_t counter) override {
-    ++counters->pause_count;
+    std::lock_guard<std::mutex> lg(mutex_);
+    if (counter > backpressure_counter_) {
+      backpressure_counter_ = counter;
+      if (!paused) ++counters->pause_count;
+      paused = true;
+    }
     inputs()[0]->PauseProducing(this, counter);
   }
   void ResumeProducing(ExecNode* output, int32_t counter) override {
-    ++counters->resume_count;
+    std::lock_guard<std::mutex> lg(mutex_);
+    if (counter > backpressure_counter_) {
+      backpressure_counter_ = counter;
+      if (paused) ++counters->resume_count;
+      paused = false;
+    }
     inputs()[0]->ResumeProducing(this, counter);
   }
 
   BackpressureCounters* counters;
+  std::mutex mutex_;
+  std::atomic<int32_t> backpressure_counter_{0};
+  bool paused{false};
 };
 
 AsyncGenerator<std::optional<ExecBatch>> GetGen(
@@ -1553,6 +1567,145 @@ TEST(AsofJoinTest, BackpressureWithBatches) {
                           /*num_r0_batches=*/50, /*num_r1_batches=*/20, /*slow_r0=*/true);
 }
 
+TEST(AsofJoinTest, PauseProducingAsofJoinSource) {
+  int batch_size = 1;
+  auto make_shift = [batch_size](int num_batches, const std::shared_ptr<Schema>& schema,
+                                 int shift) {
+    return MakeIntegerBatches(
+        {[](int row) -> int64_t { return row; },
+         [num_batches](int row) -> int64_t { return row / num_batches; },
+         [shift](int row) -> int64_t { return row * 10 + shift; }},
+        schema, num_batches, batch_size);
+  };
+  auto l_schema =
+      schema({field("time", int64()), field("key", int64()), field("l_value", int64())});
+  auto r_schema =
+      schema({field("time", int64()), field("key", int64()), field("r0_value", int64())});
+
+  auto output_schema =
+      schema({field("time", int64()), field("key", int64()), field("l_value", int64()),
+              field("key", int64()), field("r0_value", int64())});
+
+  ASSERT_OK_AND_ASSIGN(auto out_batch,
+                       MakeIntegerBatches({[](int row) -> int64_t { return row; },
+                                           [](int row) -> int64_t { return row; },
+                                           [](int row) -> int64_t { return row / 20; },
+                                           [](int row) -> int64_t { return row / 20; },
+                                           [](int row) -> int64_t { return row * 10; }},
+                                          output_schema, 20, batch_size))
+
+  ASSERT_OK_AND_ASSIGN(auto l_batches, make_shift(50, l_schema, 2));
+  ASSERT_OK_AND_ASSIGN(auto r0_batches, make_shift(50, r_schema, 1));
+  std::optional<ExecBatch> out = out_batch.batches[0];
+
+  constexpr uint32_t thresholdOfBackpressureAsof = 8;
+
+  EXPECT_OK_AND_ASSIGN(std::shared_ptr<ExecPlan> plan, ExecPlan::Make());
+  PushGenerator<std::optional<ExecBatch>> batch_producer_left;
+  PushGenerator<std::optional<ExecBatch>> batch_producer_right;
+
+  AsyncGenerator<std::optional<ExecBatch>> sink_gen;
+  BackpressureMonitor* backpressure_monitor;
+  BackpressureOptions backpressure_options(1, 2);
+  std::shared_ptr<Schema> schema_ = schema({field("data", uint32())});
+
+  BackpressureCountingNode::Register();
+
+  Declaration left{"source", SourceNodeOptions(l_schema, batch_producer_left)};
+  Declaration right{"source", SourceNodeOptions(r_schema, batch_producer_right)};
+  AsofJoinNodeOptions asof_join_opts({{{"time"}, {}}, {{"time"}, {}}}, 0);
+
+  BackpressureCounters bp_countersl, bp_countersr;
+  BackpressureCountingNode::Register();
+
+  Declaration left_count{"backpressure_count",
+                         {std::move(left)},
+                         BackpressureCountingNodeOptions(&bp_countersl)};
+
+  Declaration right_count{"backpressure_count",
+                          {std::move(right)},
+                          BackpressureCountingNodeOptions(&bp_countersr)};
+
+  Declaration asof_join{"asofjoin",
+                        {std::move(left_count), std::move(right_count)},
+                        std::move(asof_join_opts)};
+
+  ARROW_EXPECT_OK(
+      acero::Declaration::Sequence(
+          {
+              std::move(asof_join),
+              {"sink", SinkNodeOptions{&sink_gen, /*schema=*/nullptr,
+                                       backpressure_options, &backpressure_monitor}},
+          })
+          .AddToPlan(plan.get()));
+
+  ASSERT_TRUE(backpressure_monitor);
+  plan->StartProducing();
+  auto fut = plan->finished();
+
+  EXPECT_FALSE(backpressure_monitor->is_paused());
+
+  auto is_l_paused = [&]() {
+    return bp_countersl.pause_count != bp_countersl.resume_count;
+  };
+  auto is_r_paused = [&]() {
+    return bp_countersr.pause_count != bp_countersr.resume_count;
+  };
+
+  // Should be able to push kPauseIfAbove batches without triggering back pressure
+  int64_t l_cnt = 0;
+  int64_t r_cnt = 0;
+
+  EXPECT_FALSE(is_l_paused());
+  EXPECT_FALSE(is_r_paused());
+  EXPECT_FALSE(backpressure_monitor->is_paused());
+  batch_producer_left.producer().Push(l_batches.batches[l_cnt++]);
+  batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]);
+
+  // this should trigger pause on sink
+  BusyWait(3.0, [&]() { return backpressure_monitor->is_paused(); });
+  arrow::io::internal::GetIOThreadPool()->WaitForIdle();
+  arrow::internal::GetCpuThreadPool()->WaitForIdle();
+
+  // Fill up the inputs of the asof join node
+  for (uint32_t i = 0; i < thresholdOfBackpressureAsof; i++) {
+    SleepABit();
+    EXPECT_FALSE(is_l_paused());
+    EXPECT_FALSE(is_r_paused());
+    EXPECT_TRUE(backpressure_monitor->is_paused());
+    batch_producer_left.producer().Push(l_batches.batches[l_cnt++]);
+    batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]);
+  }
+
+  BusyWait(3.0, is_l_paused);
+  BusyWait(3.0, is_r_paused);
+  arrow::io::internal::GetIOThreadPool()->WaitForIdle();
+  arrow::internal::GetCpuThreadPool()->WaitForIdle();
+  // Verify pause propagates
+  EXPECT_TRUE(is_l_paused());
+  EXPECT_TRUE(is_r_paused());
+
+  std::optional<ExecBatch> opt_batch;
+
+  do {
+    ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen());
+    ASSERT_TRUE(opt_batch);
+    l_cnt -= opt_batch->length;
+  } while (l_cnt);
+
+  BusyWait(3.0, [&]() { return !is_l_paused(); });
+  BusyWait(3.0, [&]() { return !is_r_paused(); });
+  arrow::io::internal::GetIOThreadPool()->WaitForIdle();
+  arrow::internal::GetCpuThreadPool()->WaitForIdle();
+  EXPECT_FALSE(is_l_paused());
+  EXPECT_FALSE(is_r_paused());
+  EXPECT_FALSE(backpressure_monitor->is_paused());
+
+  batch_producer_left.producer().Push(IterationEnd<std::optional<ExecBatch>>());
+  batch_producer_right.producer().Push(IterationEnd<std::optional<ExecBatch>>());
+
+  ASSERT_THAT(fut, Finishes(Ok()));
+}
 template <typename BatchesMaker>
 void TestSequencing(BatchesMaker maker, int num_batches, int batch_size) {
   auto l_schema =