Merge branch 'BackpressureCombiner' into combined2

Author: Rafał Hibner

cpp/src/arrow/acero/backpressure.cc

@@ -30,15 +30,13 @@ void BackpressureController::Resume() {
 }
 
 BackpressureCombiner::BackpressureCombiner(
-    std::unique_ptr<BackpressureControl> backpressure_control)
-    : backpressure_control_(std::move(backpressure_control)) {}
+    std::unique_ptr<BackpressureControl> backpressure_control, bool pause_on_any)
+    : pause_on_any_(pause_on_any),
+      backpressure_control_(std::move(backpressure_control)) {}
 
 // Called from Source nodes
-void BackpressureCombiner::Pause(Source* output, bool strong_connection) {
+void BackpressureCombiner::Pause(Source* output) {
   std::lock_guard<std::mutex> lg(mutex_);
-  auto& paused_ = strong_connection ? strong_paused_ : weak_paused_;
-  auto& paused_count_ = strong_connection ? strong_paused_count_ : weak_paused_count_;
-
   if (!paused_[output]) {
     paused_[output] = true;
     paused_count_++;
@@ -47,25 +45,23 @@ void BackpressureCombiner::Pause(Source* output, bool strong_connection) {
 }
 
 // Called from Source nodes
-void BackpressureCombiner::Resume(Source* output, bool strong_connection) {
+void BackpressureCombiner::Resume(Source* output) {
   std::lock_guard<std::mutex> lg(mutex_);
-  auto& paused_ = strong_connection ? strong_paused_ : weak_paused_;
-  auto& paused_count_ = strong_connection ? strong_paused_count_ : weak_paused_count_;
   if (paused_.find(output) == paused_.end()) {
     paused_[output] = false;
     UpdatePauseStateUnlocked();
-  }
-  if (paused_[output]) {
+  } else if (paused_[output]) {
     paused_[output] = false;
     paused_count_--;
     UpdatePauseStateUnlocked();
   }
 }
 
 void BackpressureCombiner::UpdatePauseStateUnlocked() {
-  bool should_be_paused =
-      strong_paused_count_ > 0 ||
-      (weak_paused_count_ > 0 && weak_paused_count_ == weak_paused_.size());
+  bool should_be_paused = (paused_count_ > 0);
+  if (!pause_on_any_) {
+    should_be_paused = should_be_paused && (paused_count_ == paused_.size());
+  }
   if (should_be_paused) {
     if (!paused) {
       backpressure_control_->Pause();
@@ -79,25 +75,24 @@ void BackpressureCombiner::UpdatePauseStateUnlocked() {
   }
 }
 
-BackpressureCombiner::Source::Source(BackpressureCombiner* ctrl, bool strong_connection) {
+BackpressureCombiner::Source::Source(BackpressureCombiner* ctrl) {
   if (ctrl) {
-    AddController(ctrl, strong_connection);
+    AddController(ctrl);
   }
 }
 
-void BackpressureCombiner::Source::AddController(BackpressureCombiner* ctrl,
-                                                 bool strong_connection) {
-  ctrl->Resume(this, strong_connection);  // populate map in controller
-  connections_.push_back(Connection{ctrl, strong_connection});
+void BackpressureCombiner::Source::AddController(BackpressureCombiner* ctrl) {
+  ctrl->Resume(this);  // populate map in controller
+  connections_.push_back(ctrl);
 }
 void BackpressureCombiner::Source::Pause() {
   for (auto& conn_ : connections_) {
-    conn_.ctrl->Pause(this, conn_.strong);
+    conn_->Pause(this);
   }
 }
 void BackpressureCombiner::Source::Resume() {
   for (auto& conn_ : connections_) {
-    conn_.ctrl->Resume(this, conn_.strong);
+    conn_->Resume(this);
   }
 }
 

cpp/src/arrow/acero/backpressure.h

@@ -53,40 +53,35 @@ class BackpressureControlWrapper : public BackpressureControl {
 // strong Source within controller
 class ARROW_ACERO_EXPORT BackpressureCombiner {
  public:
-  explicit BackpressureCombiner(
-      std::unique_ptr<BackpressureControl> backpressure_control);
+  explicit BackpressureCombiner(std::unique_ptr<BackpressureControl> backpressure_control,
+                                bool pause_on_any = true);
 
   // Instances of Source can be used as usual BackpresureControl.
   // Source can be connected with one or more BackpressureCombiner
   class ARROW_ACERO_EXPORT Source : public BackpressureControl {
    public:
     // strong - strong_connection=true
     // weak - strong_connection=false
-    explicit Source(BackpressureCombiner* ctrl = nullptr, bool strong_connection = true);
-    void AddController(BackpressureCombiner* ctrl, bool strong_connection = true);
+    explicit Source(BackpressureCombiner* ctrl = nullptr);
+    void AddController(BackpressureCombiner* ctrl);
     void Pause() override;
     void Resume() override;
 
    private:
-    struct Connection {
-      BackpressureCombiner* ctrl;
-      bool strong;
-    };
-    std::vector<Connection> connections_;
+    std::vector<BackpressureCombiner*> connections_;
   };
 
  private:
   friend class Source;
-  void Pause(Source* output, bool strong_connection);
-  void Resume(Source* output, bool strong_connection);
+  void Pause(Source* output);
+  void Resume(Source* output);
 
   void UpdatePauseStateUnlocked();
+  bool pause_on_any_;
   std::unique_ptr<BackpressureControl> backpressure_control_;
   std::mutex mutex_;
-  std::unordered_map<Source*, bool> strong_paused_;
-  std::unordered_map<Source*, bool> weak_paused_;
-  size_t strong_paused_count_{0};
-  size_t weak_paused_count_{0};
+  std::unordered_map<Source*, bool> paused_;
+  size_t paused_count_{0};
   bool paused{false};
 };
 

cpp/src/arrow/acero/backpressure_test.cc

@@ -32,14 +32,17 @@ class MonitorBackpressureControl : public acero::BackpressureControl {
 
 TEST(BackpressureCombiner, Basic) {
   std::atomic<bool> paused{false};
-  BackpressureCombiner combiner(std::make_unique<MonitorBackpressureControl>(paused));
+  BackpressureCombiner or_combiner(std::make_unique<MonitorBackpressureControl>(paused));
+  BackpressureCombiner::Source strong_source1(&or_combiner);
+  BackpressureCombiner::Source strong_source2;
+  strong_source2.AddController(&or_combiner);
 
-  BackpressureCombiner::Source weak_source1(&combiner, false);
+  BackpressureCombiner and_combiner(
+      std::make_unique<BackpressureCombiner::Source>(&or_combiner),
+      /*pause_on_any=*/false);
+  BackpressureCombiner::Source weak_source1(&and_combiner);
   BackpressureCombiner::Source weak_source2;
-  weak_source2.AddController(&combiner, false);
-  BackpressureCombiner::Source strong_source1(&combiner);
-  BackpressureCombiner::Source strong_source2;
-  strong_source2.AddController(&combiner);
+  weak_source2.AddController(&and_combiner);
 
   // Any strong causes pause
   ASSERT_FALSE(paused);
@@ -102,11 +105,12 @@ TEST(BackpressureCombiner, OnlyStrong) {
 
 TEST(BackpressureCombiner, OnlyWeak) {
   std::atomic<bool> paused{false};
-  BackpressureCombiner combiner(std::make_unique<MonitorBackpressureControl>(paused));
+  BackpressureCombiner combiner(std::make_unique<MonitorBackpressureControl>(paused),
+                                false);
 
-  BackpressureCombiner::Source weak_source1(&combiner, false);
+  BackpressureCombiner::Source weak_source1(&combiner);
   BackpressureCombiner::Source weak_source2;
-  weak_source2.AddController(&combiner, false);
+  weak_source2.AddController(&combiner);
 
   // All weak cause pause
   ASSERT_FALSE(paused);