From cb14b540083480d84f09dacacc20e8df14550c26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Mon, 14 Apr 2025 13:50:59 +0000 Subject: [PATCH 01/17] Handle PauseProducing in asof_join --- cpp/src/arrow/acero/asof_join_node.cc | 53 ++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index c21af3da84f2..497044752232 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1108,6 +1108,7 @@ class AsofJoinNode : public ExecNode { EndFromProcessThread(); return; } + backpressure_future_.Wait(); if (!Process()) { return; } @@ -1520,8 +1521,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 @@ -1549,6 +1593,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_; From 7279c65eb37b02ddf047edb42b7f3c7184a27f14 Mon Sep 17 00:00:00 2001 From: kamilt Date: Mon, 19 May 2025 11:14:45 +0000 Subject: [PATCH 02/17] PauseProducing in asof-join test --- cpp/src/arrow/acero/asof_join_node_test.cc | 143 +++++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 271ad6018f2b..793295873cf8 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1552,6 +1552,149 @@ 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 thresholdOfBackpressure = 8; + constexpr uint32_t kPauseIfAbove = 4; + constexpr uint32_t kResumeIfBelow = 2; + uint32_t pause_if_above_bytes = + kPauseIfAbove * static_cast(out->TotalBufferSize()); + uint32_t resume_if_below_bytes = + kResumeIfBelow * static_cast(out->TotalBufferSize()); + 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(resume_if_below_bytes, pause_if_above_bytes); + 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"}, {}}}, 1); + + 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(); + + ASSERT_FALSE(backpressure_monitor->is_paused()); + + auto has_bp_been_applied = [&] { + for (size_t i = 0; i < 2; i++) { + const auto& counters = (i == 0) ? bp_countersl : bp_countersr; + if (counters.pause_count > 0) return true; + } + return false; + }; + + // Should be able to push kPauseIfAbove batches without triggering back pressure + uint32_t cnt = 0; + for (uint32_t i = 0; i < kPauseIfAbove; i++) { + batch_producer_left.producer().Push(l_batches.batches[i]); + batch_producer_right.producer().Push(r0_batches.batches[i]); + cnt = i; + } + cnt++; + + SleepABit(); + ASSERT_FALSE(backpressure_monitor->is_paused()); + + // One more batch should trigger back pressure + batch_producer_right.producer().Push(r0_batches.batches[cnt]); + batch_producer_left.producer().Push(l_batches.batches[cnt]); + + BusyWait(10, [&] { return backpressure_monitor->is_paused(); }); + ASSERT_TRUE(backpressure_monitor->is_paused()); + + // Fill up the inputs of the asof join node + cnt++; + for (uint32_t i = cnt; i < thresholdOfBackpressure + cnt; i++) { + batch_producer_left.producer().Push(l_batches.batches[i]); + batch_producer_right.producer().Push(r0_batches.batches[i]); + } + + BusyWait(20.0, has_bp_been_applied); + ASSERT_TRUE(has_bp_been_applied()); + + // Reading as much as we can while keeping it paused + for (uint32_t i = kPauseIfAbove; i >= kResumeIfBelow; i--) { + ASSERT_FINISHES_OK(sink_gen()); + } + SleepABit(); + ASSERT_TRUE(backpressure_monitor->is_paused()); + + // Reading one more item should open up backpressure + ASSERT_FINISHES_OK(sink_gen()); + BusyWait(10, [&] { return !backpressure_monitor->is_paused(); }); + ASSERT_FALSE(backpressure_monitor->is_paused()); + + // Cleanup + batch_producer_left.producer().Push(IterationEnd>()); + batch_producer_right.producer().Push(IterationEnd>()); + + plan->StopProducing(); + + ASSERT_TRUE(fut.Wait(kDefaultAssertFinishesWaitSeconds)); + if (!fut.status().ok()) { + ASSERT_TRUE(fut.status().IsCancelled()); + } +} template void TestSequencing(BatchesMaker maker, int num_batches, int batch_size) { auto l_schema = From b975d440558e5cc86a08e46e72bb0b86cf91e0e3 Mon Sep 17 00:00:00 2001 From: kamilt Date: Mon, 26 May 2025 08:34:02 +0000 Subject: [PATCH 03/17] Consume sink batches and clear input --- cpp/src/arrow/acero/asof_join_node_test.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 793295873cf8..5bd84682949c 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1683,8 +1683,12 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { ASSERT_FINISHES_OK(sink_gen()); BusyWait(10, [&] { return !backpressure_monitor->is_paused(); }); ASSERT_FALSE(backpressure_monitor->is_paused()); - + // Cleanup + for (size_t i = 0; i < cnt - 3; i++) { + ASSERT_FINISHES_OK(sink_gen()); + } + batch_producer_left.producer().Push(IterationEnd>()); batch_producer_right.producer().Push(IterationEnd>()); From a4da91dd693daf020a33e9a2afdf3795c16be2ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Tue, 27 May 2025 10:23:56 +0200 Subject: [PATCH 04/17] lint --- cpp/src/arrow/acero/asof_join_node_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 5bd84682949c..ba6bb598ff70 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1683,7 +1683,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { ASSERT_FINISHES_OK(sink_gen()); BusyWait(10, [&] { return !backpressure_monitor->is_paused(); }); ASSERT_FALSE(backpressure_monitor->is_paused()); - + // Cleanup for (size_t i = 0; i < cnt - 3; i++) { ASSERT_FINISHES_OK(sink_gen()); From 639af09048f2556efc07ba3b34761ea5e5cce5b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Tue, 27 May 2025 10:59:02 +0200 Subject: [PATCH 05/17] Fix UBSAN error --- cpp/src/arrow/acero/asof_join_node_test.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index ba6bb598ff70..1b4dda7d24ca 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1692,7 +1692,9 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { batch_producer_left.producer().Push(IterationEnd>()); batch_producer_right.producer().Push(IterationEnd>()); - plan->StopProducing(); + //finish gracefully + ASSERT_TRUE(batch_producer_left.producer().Close()); + ASSERT_TRUE(batch_producer_right.producer().Close()); ASSERT_TRUE(fut.Wait(kDefaultAssertFinishesWaitSeconds)); if (!fut.status().ok()) { From 5a7ae6c87d83bdd6d3c1c50abf583ed5093dd3eb Mon Sep 17 00:00:00 2001 From: kamilt Date: Wed, 28 May 2025 14:19:38 +0000 Subject: [PATCH 06/17] Corrected finishing plan --- cpp/src/arrow/acero/asof_join_node_test.cc | 109 ++++++++++++--------- 1 file changed, 64 insertions(+), 45 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 5bd84682949c..5af94a44c687 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1583,7 +1583,9 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { ASSERT_OK_AND_ASSIGN(auto r0_batches, make_shift(50, r_schema, 1)); std::optional out = out_batch.batches[0]; - constexpr uint32_t thresholdOfBackpressure = 8; + constexpr uint32_t thresholdOfBackpressureAsof = 8; + constexpr uint32_t thresholdOfBackpressureAsofLow = 4; + constexpr uint32_t kPauseIfAbove = 4; constexpr uint32_t kResumeIfBelow = 2; uint32_t pause_if_above_bytes = @@ -1608,6 +1610,12 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { BackpressureCounters bp_countersl, bp_countersr; BackpressureCountingNode::Register(); + auto wait = [](uint32_t miliseconds) { + for (size_t i = 0; i < miliseconds; i++) { + SleepABit(); + } + }; + Declaration left_count{"backpressure_count", {std::move(left)}, BackpressureCountingNodeOptions(&bp_countersl)}; @@ -1633,71 +1641,82 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { plan->StartProducing(); auto fut = plan->finished(); - ASSERT_FALSE(backpressure_monitor->is_paused()); + EXPECT_FALSE(backpressure_monitor->is_paused()); - auto has_bp_been_applied = [&] { - for (size_t i = 0; i < 2; i++) { - const auto& counters = (i == 0) ? bp_countersl : bp_countersr; - if (counters.pause_count > 0) return true; - } - return false; + 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 - uint32_t cnt = 0; + uint32_t l_cnt = 0; + uint32_t r_cnt = 0; for (uint32_t i = 0; i < kPauseIfAbove; i++) { - batch_producer_left.producer().Push(l_batches.batches[i]); - batch_producer_right.producer().Push(r0_batches.batches[i]); - cnt = i; + wait(10); + 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++]); } - cnt++; - - SleepABit(); - ASSERT_FALSE(backpressure_monitor->is_paused()); + wait(10); + EXPECT_FALSE(backpressure_monitor->is_paused()); // One more batch should trigger back pressure - batch_producer_right.producer().Push(r0_batches.batches[cnt]); - batch_producer_left.producer().Push(l_batches.batches[cnt]); + batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); + batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); - BusyWait(10, [&] { return backpressure_monitor->is_paused(); }); - ASSERT_TRUE(backpressure_monitor->is_paused()); + wait(10); + EXPECT_TRUE(backpressure_monitor->is_paused()); // Fill up the inputs of the asof join node - cnt++; - for (uint32_t i = cnt; i < thresholdOfBackpressure + cnt; i++) { - batch_producer_left.producer().Push(l_batches.batches[i]); - batch_producer_right.producer().Push(r0_batches.batches[i]); + for (uint32_t i = 0; i < thresholdOfBackpressureAsof; i++) { + wait(10); + EXPECT_FALSE(is_l_paused()); + EXPECT_FALSE(is_r_paused()); + batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); + batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); } - BusyWait(20.0, has_bp_been_applied); - ASSERT_TRUE(has_bp_been_applied()); + std::optional opt_batch; - // Reading as much as we can while keeping it paused - for (uint32_t i = kPauseIfAbove; i >= kResumeIfBelow; i--) { - ASSERT_FINISHES_OK(sink_gen()); + // Read the batches from the sink to open up input of the asof join node + for (uint32_t i = 0; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; + i++) { + wait(10); + EXPECT_TRUE(is_l_paused()); + EXPECT_TRUE(is_r_paused()); + EXPECT_TRUE(backpressure_monitor->is_paused()); + + ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); + EXPECT_TRUE(opt_batch); } - SleepABit(); - ASSERT_TRUE(backpressure_monitor->is_paused()); - - // Reading one more item should open up backpressure - ASSERT_FINISHES_OK(sink_gen()); - BusyWait(10, [&] { return !backpressure_monitor->is_paused(); }); - ASSERT_FALSE(backpressure_monitor->is_paused()); - - // Cleanup - for (size_t i = 0; i < cnt - 3; i++) { - ASSERT_FINISHES_OK(sink_gen()); + + // Finish the batches in the left and right producers + for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + kResumeIfBelow + 2; i++) { + wait(10); + EXPECT_FALSE(is_l_paused()); + EXPECT_FALSE(is_r_paused()); + EXPECT_TRUE(backpressure_monitor->is_paused()); + + ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); + EXPECT_TRUE(opt_batch); } + wait(10); + EXPECT_FALSE(is_l_paused()); + EXPECT_FALSE(is_r_paused()); + EXPECT_FALSE(backpressure_monitor->is_paused()); + + ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); + EXPECT_TRUE(opt_batch); batch_producer_left.producer().Push(IterationEnd>()); batch_producer_right.producer().Push(IterationEnd>()); - plan->StopProducing(); + ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); + EXPECT_FALSE(opt_batch); - ASSERT_TRUE(fut.Wait(kDefaultAssertFinishesWaitSeconds)); - if (!fut.status().ok()) { - ASSERT_TRUE(fut.status().IsCancelled()); - } + ASSERT_THAT(fut, Finishes(Ok())); } template void TestSequencing(BatchesMaker maker, int num_batches, int batch_size) { From 23cb51954100075fc738983782a47888f49c2157 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 28 May 2025 15:13:14 +0000 Subject: [PATCH 07/17] Busywait for completion --- cpp/src/arrow/acero/asof_join_node_test.cc | 25 +++++++++++----------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 5af94a44c687..4c443c37ff29 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1610,12 +1610,6 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { BackpressureCounters bp_countersl, bp_countersr; BackpressureCountingNode::Register(); - auto wait = [](uint32_t miliseconds) { - for (size_t i = 0; i < miliseconds; i++) { - SleepABit(); - } - }; - Declaration left_count{"backpressure_count", {std::move(left)}, BackpressureCountingNodeOptions(&bp_countersl)}; @@ -1654,24 +1648,24 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { uint32_t l_cnt = 0; uint32_t r_cnt = 0; for (uint32_t i = 0; i < kPauseIfAbove; i++) { - wait(10); + SleepABit(); 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++]); } - wait(10); + SleepABit(); EXPECT_FALSE(backpressure_monitor->is_paused()); // One more batch should trigger back pressure batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); - wait(10); + SleepABit(); EXPECT_TRUE(backpressure_monitor->is_paused()); // Fill up the inputs of the asof join node for (uint32_t i = 0; i < thresholdOfBackpressureAsof; i++) { - wait(10); + SleepABit(); EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); @@ -1683,7 +1677,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { // Read the batches from the sink to open up input of the asof join node for (uint32_t i = 0; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; i++) { - wait(10); + SleepABit(); EXPECT_TRUE(is_l_paused()); EXPECT_TRUE(is_r_paused()); EXPECT_TRUE(backpressure_monitor->is_paused()); @@ -1692,9 +1686,12 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { EXPECT_TRUE(opt_batch); } + BusyWait(5.0, [&]() { return !is_l_paused(); }); + BusyWait(5.0, [&]() { return !is_r_paused(); }); + // Finish the batches in the left and right producers for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + kResumeIfBelow + 2; i++) { - wait(10); + SleepABit(); EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); EXPECT_TRUE(backpressure_monitor->is_paused()); @@ -1702,7 +1699,9 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); EXPECT_TRUE(opt_batch); } - wait(10); + + BusyWait(5.0, [&]() { return !backpressure_monitor->is_paused(); }); + EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); EXPECT_FALSE(backpressure_monitor->is_paused()); From 5e9de2b0abdc2385fcad1ce1895a5dda3629403f Mon Sep 17 00:00:00 2001 From: kamilt Date: Fri, 30 May 2025 09:28:04 +0000 Subject: [PATCH 08/17] Add BusyWait for state change --- cpp/src/arrow/acero/asof_join_node_test.cc | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 4c443c37ff29..9ef063316405 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1660,7 +1660,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); - SleepABit(); + BusyWait(5, [&] { return backpressure_monitor->is_paused(); }); EXPECT_TRUE(backpressure_monitor->is_paused()); // Fill up the inputs of the asof join node @@ -1673,7 +1673,8 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { } std::optional opt_batch; - + BusyWait(5.0, [&]() { return is_l_paused(); }); + BusyWait(5.0, [&]() { return is_r_paused(); }); // Read the batches from the sink to open up input of the asof join node for (uint32_t i = 0; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; i++) { @@ -1690,25 +1691,18 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { BusyWait(5.0, [&]() { return !is_r_paused(); }); // Finish the batches in the left and right producers - for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + kResumeIfBelow + 2; i++) { + for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + kResumeIfBelow + 3; i++) { SleepABit(); EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); - EXPECT_TRUE(backpressure_monitor->is_paused()); - ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); EXPECT_TRUE(opt_batch); } - BusyWait(5.0, [&]() { return !backpressure_monitor->is_paused(); }); - EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); EXPECT_FALSE(backpressure_monitor->is_paused()); - ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); - EXPECT_TRUE(opt_batch); - batch_producer_left.producer().Push(IterationEnd>()); batch_producer_right.producer().Push(IterationEnd>()); From 5109d5b8be2c4989e1b5f5b562b8d634d6fca0cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Mon, 2 Jun 2025 14:49:12 +0000 Subject: [PATCH 09/17] Fix asof_join pause --- cpp/src/arrow/acero/asof_join_node.cc | 2 +- cpp/src/arrow/acero/asof_join_node_test.cc | 51 ++++++++++++---------- 2 files changed, 28 insertions(+), 25 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 497044752232..65c1a0b403ea 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1069,6 +1069,7 @@ class AsofJoinNode : public ExecNode { // Process batches while we have data for (;;) { + backpressure_future_.Wait(); Result> result = ProcessInner(); if (result.ok()) { @@ -1108,7 +1109,6 @@ class AsofJoinNode : public ExecNode { EndFromProcessThread(); return; } - backpressure_future_.Wait(); if (!Process()) { return; } diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 9ef063316405..eac2ac86a124 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/gtest_util.h" #include "arrow/testing/matchers.h" #include "arrow/testing/random.h" @@ -1586,19 +1587,13 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { constexpr uint32_t thresholdOfBackpressureAsof = 8; constexpr uint32_t thresholdOfBackpressureAsofLow = 4; - constexpr uint32_t kPauseIfAbove = 4; - constexpr uint32_t kResumeIfBelow = 2; - uint32_t pause_if_above_bytes = - kPauseIfAbove * static_cast(out->TotalBufferSize()); - uint32_t resume_if_below_bytes = - kResumeIfBelow * static_cast(out->TotalBufferSize()); 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(resume_if_below_bytes, pause_if_above_bytes); + BackpressureOptions backpressure_options(1, 2); std::shared_ptr schema_ = schema({field("data", uint32())}); BackpressureCountingNode::Register(); @@ -1647,34 +1642,43 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { // Should be able to push kPauseIfAbove batches without triggering back pressure uint32_t l_cnt = 0; uint32_t r_cnt = 0; - for (uint32_t i = 0; i < kPauseIfAbove; i++) { - SleepABit(); - 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++]); - } - - SleepABit(); + // for (uint32_t i = 0; i < 1; i++) { + EXPECT_FALSE(is_l_paused()); + EXPECT_FALSE(is_r_paused()); EXPECT_FALSE(backpressure_monitor->is_paused()); - // One more batch should trigger back pressure + batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); + // } + + // One more batch should trigger back pressure + + BusyWait(60.0, [&]() { return backpressure_monitor->is_paused(); }); + 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(); - BusyWait(5, [&] { return backpressure_monitor->is_paused(); }); + 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 = 0; i < thresholdOfBackpressureAsof; i++) { + for (uint32_t i = 1; i < thresholdOfBackpressureAsof; i++) { SleepABit(); - EXPECT_FALSE(is_l_paused()); + EXPECT_FALSE(is_l_paused()) << i; EXPECT_FALSE(is_r_paused()); batch_producer_left.producer().Push(l_batches.batches[l_cnt++]); batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); } std::optional opt_batch; - BusyWait(5.0, [&]() { return is_l_paused(); }); - BusyWait(5.0, [&]() { return is_r_paused(); }); + arrow::internal::GetCpuThreadPool()->WaitForIdle(); // Read the batches from the sink to open up input of the asof join node for (uint32_t i = 0; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; i++) { @@ -1687,11 +1691,10 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { EXPECT_TRUE(opt_batch); } - BusyWait(5.0, [&]() { return !is_l_paused(); }); - BusyWait(5.0, [&]() { return !is_r_paused(); }); + arrow::internal::GetCpuThreadPool()->WaitForIdle(); // Finish the batches in the left and right producers - for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + kResumeIfBelow + 3; i++) { + for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + 1; i++) { SleepABit(); EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); From b60041a6aa82364a3ddf69704778f2c40d620c4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Mon, 2 Jun 2025 15:36:12 +0000 Subject: [PATCH 10/17] Fix test --- cpp/src/arrow/acero/asof_join_node_test.cc | 42 ++++++++++------------ 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 19dee4049fd8..ca9cfd6e0d62 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1532,7 +1532,7 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches, return true; }; - BusyWait(60.0, has_bp_been_applied); + BusyWait(3.0, has_bp_been_applied); ASSERT_TRUE(has_bp_been_applied()); gate.ReleaseAllBatches(); @@ -1649,11 +1649,10 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { 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++]); - // } - // One more batch should trigger back pressure + // this should trigger pause on sink - BusyWait(60.0, [&]() { return backpressure_monitor->is_paused(); }); + BusyWait(3.0, [&]() { return backpressure_monitor->is_paused(); }); arrow::io::internal::GetIOThreadPool()->WaitForIdle(); arrow::internal::GetCpuThreadPool()->WaitForIdle(); @@ -1678,38 +1677,33 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { batch_producer_right.producer().Push(r0_batches.batches[r_cnt++]); } - std::optional opt_batch; + BusyWait(3.0, is_l_paused); + BusyWait(3.0, is_r_paused); + arrow::io::internal::GetIOThreadPool()->WaitForIdle(); arrow::internal::GetCpuThreadPool()->WaitForIdle(); - // Read the batches from the sink to open up input of the asof join node - for (uint32_t i = 0; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; - i++) { - SleepABit(); - EXPECT_TRUE(is_l_paused()); - EXPECT_TRUE(is_r_paused()); - EXPECT_TRUE(backpressure_monitor->is_paused()); + // Verify pause propagates + EXPECT_TRUE(is_l_paused()); + EXPECT_TRUE(is_r_paused()); - ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); - EXPECT_TRUE(opt_batch); - } + batch_producer_left.producer().Push(IterationEnd>()); + batch_producer_right.producer().Push(IterationEnd>()); - arrow::internal::GetCpuThreadPool()->WaitForIdle(); + std::optional opt_batch; - // Finish the batches in the left and right producers - for (uint32_t i = 0; i < thresholdOfBackpressureAsofLow + 1; i++) { - SleepABit(); - EXPECT_FALSE(is_l_paused()); - EXPECT_FALSE(is_r_paused()); + for (uint32_t i = 1; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; + i++) { ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); EXPECT_TRUE(opt_batch); } + 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_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); EXPECT_FALSE(opt_batch); From 59008c2390b871337c3c6914351c1bcd0edb527b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Mon, 2 Jun 2025 16:07:41 +0000 Subject: [PATCH 11/17] Revert spurious change --- cpp/src/arrow/acero/asof_join_node_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index ca9cfd6e0d62..60863d92e325 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1532,7 +1532,7 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches, return true; }; - BusyWait(3.0, has_bp_been_applied); + BusyWait(60.0, has_bp_been_applied); ASSERT_TRUE(has_bp_been_applied()); gate.ReleaseAllBatches(); From 3f1dc3c3a64bdde1699db43ccdd23a275631d537 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Tue, 3 Jun 2025 12:56:39 +0000 Subject: [PATCH 12/17] Fix test --- cpp/src/arrow/acero/asof_join_node_test.cc | 41 +++++++++++++--------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 60863d92e325..9b4168218098 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1413,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( @@ -1586,7 +1599,6 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { std::optional out = out_batch.batches[0]; constexpr uint32_t thresholdOfBackpressureAsof = 8; - constexpr uint32_t thresholdOfBackpressureAsofLow = 4; EXPECT_OK_AND_ASSIGN(std::shared_ptr plan, ExecPlan::Make()); PushGenerator> batch_producer_left; @@ -1601,7 +1613,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { Declaration left{"source", SourceNodeOptions(l_schema, batch_producer_left)}; Declaration right{"source", SourceNodeOptions(r_schema, batch_producer_right)}; - AsofJoinNodeOptions asof_join_opts({{{"time"}, {}}, {{"time"}, {}}}, 1); + AsofJoinNodeOptions asof_join_opts({{{"time"}, {}}, {{"time"}, {}}}, 0); BackpressureCounters bp_countersl, bp_countersr; BackpressureCountingNode::Register(); @@ -1643,7 +1655,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { // Should be able to push kPauseIfAbove batches without triggering back pressure uint32_t l_cnt = 0; uint32_t r_cnt = 0; - // for (uint32_t i = 0; i < 1; i++) { + EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused()); EXPECT_FALSE(backpressure_monitor->is_paused()); @@ -1651,7 +1663,6 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { 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(); @@ -1659,6 +1670,7 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { 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++]); @@ -1685,27 +1697,24 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { EXPECT_TRUE(is_l_paused()); EXPECT_TRUE(is_r_paused()); - batch_producer_left.producer().Push(IterationEnd>()); - batch_producer_right.producer().Push(IterationEnd>()); - std::optional opt_batch; - for (uint32_t i = 1; i < thresholdOfBackpressureAsof - thresholdOfBackpressureAsofLow; - i++) { + do { ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); - EXPECT_TRUE(opt_batch); - } + 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()); - ASSERT_FINISHES_OK_AND_ASSIGN(opt_batch, sink_gen()); - EXPECT_FALSE(opt_batch); + batch_producer_left.producer().Push(IterationEnd>()); + batch_producer_right.producer().Push(IterationEnd>()); ASSERT_THAT(fut, Finishes(Ok())); } From 208f61d68ddae37a14a0798cd9ce54f2ef52a45b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Tue, 3 Jun 2025 15:52:01 +0000 Subject: [PATCH 13/17] Do not hold lock when backpressure is applied --- cpp/src/arrow/acero/asof_join_node.cc | 41 ++++++++++++++------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index b1a2b2d10e08..6c6d737470a7 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1060,30 +1060,31 @@ class AsofJoinNode : public ExecNode { } bool Process() { - std::lock_guard guard(gate_); - if (!CheckEnded()) { - return false; - } - // Process batches while we have data for (;;) { backpressure_future_.Wait(); - Result> 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)); + { + std::lock_guard guard(gate_); + if (!CheckEnded()) { + return false; + } + Result> 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; } - } else { - EndFromProcessThread(result.status()); - return false; } } From b98589cd4ba65a160758e5305c360b9449dfe908 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 4 Jun 2025 09:51:28 +0000 Subject: [PATCH 14/17] Cleanup test --- cpp/src/arrow/acero/asof_join_node_test.cc | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 9b4168218098..2f48d48b3bce 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1667,24 +1667,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++]); } From 51596f53957f836a59e2442577325211af92198b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 4 Jun 2025 10:01:34 +0000 Subject: [PATCH 15/17] Release lock when pushing result --- cpp/src/arrow/acero/asof_join_node.cc | 32 +++++++++++++-------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 6c6d737470a7..2eebe20f39e3 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/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> result; { std::lock_guard guard(gate_); if (!CheckEnded()) { return false; } - Result> 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; } } From d8d36ebf2b81329a4b92e18d16a3cd3984b3fc9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Fri, 4 Jul 2025 08:21:41 +0000 Subject: [PATCH 16/17] Fix concurrent access to backpressure_future_ --- cpp/src/arrow/acero/asof_join_node.cc | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 2eebe20f39e3..0976cc7c34ca 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1062,7 +1062,13 @@ class AsofJoinNode : public ExecNode { bool Process() { // Process batches while we have data for (;;) { - backpressure_future_.Wait(); + Future<> to_wait; + { + std::lock_guard lg(backpressure_mutex_); + to_wait = backpressure_future_; + } + to_wait.Wait(); + Result> result; { std::lock_guard guard(gate_); From 0a1db5252b5a9c02734ddef8c518b707236753fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Fri, 4 Jul 2025 08:37:31 +0000 Subject: [PATCH 17/17] Fix build error --- cpp/src/arrow/acero/asof_join_node_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 2f48d48b3bce..dcfc4c188772 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1653,8 +1653,8 @@ TEST(AsofJoinTest, PauseProducingAsofJoinSource) { }; // Should be able to push kPauseIfAbove batches without triggering back pressure - uint32_t l_cnt = 0; - uint32_t r_cnt = 0; + int64_t l_cnt = 0; + int64_t r_cnt = 0; EXPECT_FALSE(is_l_paused()); EXPECT_FALSE(is_r_paused());