diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 3970050e5022..5b205c612b2e 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1063,15 +1063,23 @@ class AsofJoinNode : public ExecNode { } bool Process() { - std::lock_guard guard(gate_); - if (!CheckEnded()) { - return false; - } - // Process batches while we have data for (;;) { - Result> result = ProcessInner(); + Future<> to_wait; + { + std::lock_guard lg(backpressure_mutex_); + to_wait = backpressure_future_; + } + to_wait.Wait(); + Result> result; + { + std::lock_guard 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 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 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 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> state_; std::mutex gate_; + + std::mutex backpressure_mutex_; + std::atomic last_backpressure_counter_{0}; + Future<> backpressure_future_ = Future<>::MakeFinished(); + TolType tolerance_; #ifndef NDEBUG std::ostream* debug_os_; diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 59a9b4ebba12..dcfc4c188772 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/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 ProcessBatch(ExecBatch batch) override { return batch; } void PauseProducing(ExecNode* output, int32_t counter) override { - ++counters->pause_count; + std::lock_guard 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 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 backpressure_counter_{0}; + bool paused{false}; }; AsyncGenerator> 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, + 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 out = out_batch.batches[0]; + + constexpr uint32_t thresholdOfBackpressureAsof = 8; + + EXPECT_OK_AND_ASSIGN(std::shared_ptr plan, ExecPlan::Make()); + PushGenerator> batch_producer_left; + PushGenerator> batch_producer_right; + + AsyncGenerator> sink_gen; + BackpressureMonitor* backpressure_monitor; + BackpressureOptions backpressure_options(1, 2); + std::shared_ptr 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 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>()); + batch_producer_right.producer().Push(IterationEnd>()); + + ASSERT_THAT(fut, Finishes(Ok())); +} template void TestSequencing(BatchesMaker maker, int num_batches, int batch_size) { auto l_schema =