Merge branch 'asof_join_pause' into combined2

Author: Rafał Hibner

cpp/src/arrow/acero/asof_join_node.cc

@@ -1063,28 +1063,28 @@ class AsofJoinNode : public ExecNode {
     // Process batches while we have data
     for (;;) {
       backpressure_future_.Wait();
+      Result<std::shared_ptr<RecordBatch>> result;
       {
         std::lock_guard<std::mutex> guard(gate_);
         if (!CheckEnded()) {
           return false;
         }
-        Result<std::shared_ptr<RecordBatch>> result = ProcessInner();
-
-        if (result.ok()) {
-          auto out_rb = *result;
-          if (!out_rb) break;
-          ExecBatch out_b(*out_rb);
-          out_b.index = batches_produced_++;
-          DEBUG_SYNC(this, "produce batch ", out_b.index, ":", DEBUG_MANIP(std::endl),
-                     out_rb->ToString(), DEBUG_MANIP(std::endl));
-          Status st = output_->InputReceived(this, std::move(out_b));
-          if (!st.ok()) {
-            EndFromProcessThread(std::move(st));
-          }
-        } else {
-          EndFromProcessThread(result.status());
-          return false;
+        result = ProcessInner();
+      }
+      if (result.ok()) {
+        auto out_rb = *result;
+        if (!out_rb) break;
+        ExecBatch out_b(*out_rb);
+        out_b.index = batches_produced_++;
+        DEBUG_SYNC(this, "produce batch ", out_b.index, ":", DEBUG_MANIP(std::endl),
+                   out_rb->ToString(), DEBUG_MANIP(std::endl));
+        Status st = output_->InputReceived(this, std::move(out_b));
+        if (!st.ok()) {
+          EndFromProcessThread(std::move(st));
         }
+      } else {
+        EndFromProcessThread(result.status());
+        return false;
       }
     }
 

cpp/src/arrow/acero/asof_join_node_test.cc

@@ -1659,24 +1659,12 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) {
   arrow::io::internal::GetIOThreadPool()->WaitForIdle();
   arrow::internal::GetCpuThreadPool()->WaitForIdle();
 
-  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++]);
-  arrow::internal::GetCpuThreadPool()->WaitForIdle();
-  batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]);
-  arrow::internal::GetCpuThreadPool()->WaitForIdle();
-
-  EXPECT_FALSE(is_l_paused());
-  EXPECT_FALSE(is_r_paused());
-  EXPECT_TRUE(backpressure_monitor->is_paused());
-
   // Fill up the inputs of the asof join node
-  for (uint32_t i = 1; i < thresholdOfBackpressureAsof; i++) {
+  for (uint32_t i = 0; i < thresholdOfBackpressureAsof; i++) {
     SleepABit();
-    EXPECT_FALSE(is_l_paused()) << i;
+    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++]);
   }