Update

dask · Feb 24, 2024 · 8bc45b9 · 8bc45b9
1 parent 9cbcec3
commit 8bc45b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dask_expr/_reductions.py b/dask_expr/_reductions.py
@@ -518,7 +518,7 @@ def _lower(self):
         )
 
     def _substitute_branch_id(self, branch_id):
-        if self.should_shuffle:
+        if self._reuse_consumer:
             # We are lowering into a Shuffle, so we are a consumer ourselves and
             # we have to consume the branch_id of our parents
             return super()._substitute_branch_id(branch_id)
@@ -528,7 +528,7 @@ def _reuse_down(self):
         if self._branch_id.branch_id != 0:
             return
 
-        if self.should_shuffle:
+        if self._reuse_consumer:
             # We are lowering into a Shuffle, so we are a consumer ourselves
             return