From 8bc45b92397d07dc0f88acc7ac8e463af0d23c5e Mon Sep 17 00:00:00 2001
From: Patrick Hoefler <61934744+phofl@users.noreply.github.com>
Date: Sat, 24 Feb 2024 15:55:52 +0100
Subject: [PATCH] Update

---
 dask_expr/_reductions.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dask_expr/_reductions.py b/dask_expr/_reductions.py
index 19f481b8c..b41e624ab 100644
--- 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