From e3b7610b98fd2f26e9dc4172e8db5bcf98ab0926 Mon Sep 17 00:00:00 2001
From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com>
Date: Tue, 8 Sep 2026 15:43:32 +0300
Subject: [PATCH] fix: mark partial reduce hash aggregate stream as can handle
 back pressure

---
 datafusion/physical-plan/src/aggregates/hash_stream.rs        | 3 +++
 .../physical-plan/src/aggregates/partial_reduce_stream.rs     | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs
index 7ebf32c3edfc6..2f18afae961bc 100644
--- a/datafusion/physical-plan/src/aggregates/hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs
@@ -425,6 +425,9 @@ impl PartialHashAggregateStream {
 
         let reservation =
             MemoryConsumer::new(format!("PartialHashAggregateStream[{partition}]"))
+                // We interpret 'can spill' as 'can handle memory back pressure'.
+                // This value needs to be set to true for the default memory pool implementations
+                // to ensure fair application of back pressure amongst the memory consumers.
                 .with_can_spill(true)
                 .register(context.memory_pool());
 
diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
index d8f1447bc9521..24db34c03cbd7 100644
--- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
@@ -203,6 +203,10 @@ impl PartialReduceHashAggregateStream {
 
         let reservation =
             MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]"))
+                // We interpret 'can spill' as 'can handle memory back pressure'.
+                // This value needs to be set to true for the default memory pool implementations
+                // to ensure fair application of back pressure amongst the memory consumers.
+                .with_can_spill(true)
                 .register(context.memory_pool());
 
         Ok(Self {