Merge branch 'BackpressureCombiner' into combined2

Author: Rafał Hibner

cpp/src/arrow/acero/backpressure.cc

@@ -64,7 +64,8 @@ void BackpressureCombiner::Resume(Source* output, bool strong_connection) {
 
 void BackpressureCombiner::UpdatePauseStateUnlocked() {
   bool should_be_paused =
-      strong_paused_count_ > 0 || weak_paused_count_ == weak_paused_.size();
+      strong_paused_count_ > 0 ||
+      (weak_paused_count_ > 0 && weak_paused_count_ == weak_paused_.size());
   if (should_be_paused) {
     if (!paused) {
       backpressure_control_->Pause();

cpp/src/arrow/acero/backpressure.h

@@ -22,7 +22,7 @@
 namespace arrow::acero {
 
 // Generic backpressure controller for ExecNode
-class BackpressureController : public BackpressureControl {
+class ARROW_ACERO_EXPORT BackpressureController : public BackpressureControl {
  public:
   BackpressureController(ExecNode* node, ExecNode* output,
                          std::atomic<int32_t>& backpressure_counter);
@@ -36,6 +36,18 @@ class BackpressureController : public BackpressureControl {
   std::atomic<int32_t>& backpressure_counter_;
 };
 
+template <typename T>
+class ARROW_ACERO_EXPORT BackpressureControlWrapper : public BackpressureControl {
+ public:
+  explicit BackpressureControlWrapper(T* obj) : obj_(obj) {}
+
+  void Pause() override { obj_->Pause(); }
+  void Resume() override { obj_->Resume(); }
+
+ private:
+  T* obj_;
+};
+
 // Provides infrastructure of combining multiple backpressure sources and propagate the
 // result into BackpressureControl There are two types of Source: strong - pause on any
 // strong Source within controller

cpp/src/arrow/acero/backpressure_test.cc

@@ -81,5 +81,44 @@ TEST(BackpressureCombiner, Basic) {
   ASSERT_FALSE(paused);
 }
 
+TEST(BackpressureCombiner, OnlyStrong) {
+  std::atomic<bool> paused{false};
+  BackpressureCombiner combiner(std::make_unique<MonitorBackpressureControl>(paused));
+  BackpressureCombiner::Source strong_source1(&combiner);
+  BackpressureCombiner::Source strong_source2;
+  strong_source2.AddController(&combiner);
+
+  // Any strong causes pause
+  ASSERT_FALSE(paused);
+  strong_source1.Pause();
+  ASSERT_TRUE(paused);
+  strong_source2.Pause();
+  ASSERT_TRUE(paused);
+  strong_source1.Resume();
+  ASSERT_TRUE(paused);
+  strong_source2.Resume();
+  ASSERT_FALSE(paused);
+}
+
+TEST(BackpressureCombiner, OnlyWeak) {
+  std::atomic<bool> paused{false};
+  BackpressureCombiner combiner(std::make_unique<MonitorBackpressureControl>(paused));
+
+  BackpressureCombiner::Source weak_source1(&combiner, false);
+  BackpressureCombiner::Source weak_source2;
+  weak_source2.AddController(&combiner, false);
+
+  // All weak cause pause
+  ASSERT_FALSE(paused);
+  weak_source1.Pause();
+  ASSERT_FALSE(paused);
+  weak_source2.Pause();
+  ASSERT_TRUE(paused);
+  weak_source1.Resume();
+  ASSERT_FALSE(paused);
+  weak_source2.Resume();
+  ASSERT_FALSE(paused);
+}
+
 }  // namespace acero
 }  // namespace arrow

cpp/src/arrow/acero/meson.build

@@ -51,6 +51,7 @@ arrow_acero_srcs = [
     'groupby_aggregate_node.cc',
     'aggregate_internal.cc',
     'asof_join_node.cc',
+    'backpressure.cc',
     'bloom_filter.cc',
     'exec_plan.cc',
     'fetch_node.cc',
@@ -90,6 +91,7 @@ arrow_acero_dep = declare_dependency(link_with: [arrow_acero_lib])
 arrow_acero_testing_sources = ['test_nodes.cc', 'test_util_internal.cc']
 
 arrow_acero_tests = {
+    'backpressure-test': {'sources': ['backpressure_test.cc']},
     'plan-test': {'sources': ['plan_test.cc', 'test_nodes_test.cc']},
     'source-node-test': {'sources': ['source_node_test.cc']},
     'fetch-node-test': {'sources': ['fetch_node_test.cc']},