You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/25 12:11:45 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6205-570a97cce869a7f11738d43d7b179c0c9c53be7c created (now e90084420a)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6205-570a97cce869a7f11738d43d7b179c0c9c53be7c
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at e90084420a fix: check each fair_unified reservation against its own share (#6205)
This branch includes the following new commits:
new e90084420a fix: check each fair_unified reservation against its own share (#6205)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: check each fair_unified reservation against its own share (#6205)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6205-570a97cce869a7f11738d43d7b179c0c9c53be7c
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit e90084420ab7acc4d1fd3a6db41f9ef9cb5a5087
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Sep 25 12:11:33 2026 +0000
fix: check each fair_unified reservation against its own share (#6205)
* fix: check each fair_unified reservation against its own share
Since the DataFusion 53 upgrade (#3629), CometFairMemoryPool::try_grow has
compared the pool's total reserved bytes against pool_size / num_consumers.
That caps the whole task at one consumer's share, and every new registration
tightens the ceiling on the consumers already running.
The comment that justified the change holds for shrink but not for try_grow.
MemoryReservation::try_grow calls pool.try_grow before adding to its size, so
reservation.size() is the size before the request. Check the reservation
against its share again, and separately keep the pool's total within
pool_size, since the shares alone do not bound it once a consumer registers
late or one consumer holds sibling reservations.
Update the tuning guide, the memory management guide and the memory PR review
skill, which described the pool-wide comparison.
Closes #5961.
* fix: charge sibling fair_unified reservations to their consumer's share
Checking each reservation against the share let an operator that holds
several reservations under one consumer, such as a sort's streaming merge
with one per batch, take other consumers' shares, bounded only by the pool
total. Keep a running total for each consumer id instead, so all of a
consumer's reservations count against one share. The pool no longer reads
reservation.size(), so the check also stops depending on when DataFusion
updates it.
Note in the tuning guide that tasks can reserve more before they spill
than they could in Comet 0.15.0 through 1.0.0.
---
.ai/skills/review-comet-memory-pr/SKILL.md | 19 +-
docs/source/contributor-guide/memory_management.md | 41 +++--
docs/source/user-guide/latest/tuning.md | 9 +-
.../core/src/execution/memory_pools/fair_pool.rs | 191 +++++++++++++++++----
4 files changed, 206 insertions(+), 54 deletions(-)
diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md
index 398adf7472..fbd4610b20 100644
--- a/.ai/skills/review-comet-memory-pr/SKILL.md
+++ b/.ai/skills/review-comet-memory-pr/SKILL.md
@@ -100,11 +100,17 @@ configuration, from the inside out:
- [ ] A new decorator forwards **every** `MemoryPool` method to its inner pool. A partial
implementation makes `reserved()` disagree between levels.
-- [ ] **`fair_unified` compares against the shared total, not a per-consumer quota.** It divides
- `pool_size` by the number of registered consumers and rejects if the pool's total reserved
- plus the request exceeds that quotient. With two consumers and an 8 GiB pool, once one holds
- 3 GiB a 2 GiB request from the other is refused. If the PR changes this comparison it changes
- when every query spills, so it needs benchmark evidence, not reasoning.
+- [ ] **`fair_unified` checks the requesting consumer against its share, and the total against
+ the pool.** It divides `pool_size` by the number of registered consumers and rejects if what
+ the consumer holds plus the request exceeds that quotient, or if the pool's total reserved
+ plus the request exceeds `pool_size`. With two consumers and an 8 GiB pool, once one holds
+ 3 GiB the other can still reserve up to 4 GiB. The pool keeps a running total for each
+ consumer id, so the sibling reservations that `new_empty()`, `split()` and `take()` create
+ count against one share. A PR that checks `reservation.size()` instead lets an operator with
+ several reservations, such as a sort's streaming merge, take other consumers' shares. It
+ also depends on when DataFusion updates the size, which is after it calls `try_grow` but
+ before it calls `shrink`. If the PR changes either comparison it changes when every query
+ spills, so it needs benchmark evidence, not reasoning.
- [ ] **`num_consumers` counts every consumer in the task**, across every native plan, because the
pool is task-shared.
- [ ] **Task-shared pool lifetime.** `acquire_task_shared_pool` keeps a process-wide
@@ -194,7 +200,8 @@ will not catch. Check:
- The **pool stack** diagram, if a decorator is added, removed, or reordered
- The **pool type** table, if a pool type is added, removed, or changes how it is sized
- The `memory_limit` formula, if the budget calculation changes
-- The **fair pool** description, which spells out the shared-total comparison and its consequence
+- The **fair pool** description, which spells out the share and pool-total checks and their
+ consequences
- The **Crossing the FFI boundary** section, which names the operators that reserve for imported
batches
- The **Open problems** list. If the PR closes one of these gaps, the entry must be removed or
diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 339e8f03e7..fec46087a1 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -266,21 +266,34 @@ JNI, which goes through Spark's ordinary `TaskMemoryManager`. That means:
is refused unless Spark can cover both, so operators spill until the debt is repaid. Both pools'
`Display` output and their `try_grow` errors report the current overcommit.
-`CometFairMemoryPool` additionally applies a local check before it asks Spark. It divides
-`pool_size` by the number of consumers currently registered with the pool and rejects the request if
-the pool's _total_ reserved bytes plus the request would exceed that quotient. Two details matter
-for tuning:
-
-- The comparison is against the shared total (`state.used`), not against the requesting consumer's
- own reservation. With an 8 GiB pool and two registered consumers, once one of them holds 3 GiB a
- 2 GiB request from the other is rejected, even though neither would exceed 4 GiB. In effect the
- usable pool shrinks to `pool_size / num_consumers` in aggregate as soon as more than one consumer
- is registered.
-- The pool is task-shared (see below), so `num_consumers` counts every registered consumer across
- every native plan in the task, not just the plan making the request.
+`CometFairMemoryPool` additionally applies two local checks before it asks Spark, and refuses the
+request without calling Spark if either fails:
+
+- **The requesting consumer against its share.** The share is `pool_size` divided by the number of
+ consumers currently registered with the pool. What the consumer already holds plus the request
+ must not exceed it. That counts all of the consumer's reservations, including the sibling
+ reservations that `new_empty()`, `split()` and `take()` create, so an operator that holds several
+ reservations still gets one share. A sort's streaming merge, for example, creates one for each
+ batch it reads. The pool keeps a running total for each consumer, because `reservation.size()`
+ covers only one reservation. With an 8 GiB pool and two registered consumers, each share is
+ 4 GiB, so once one consumer holds 3 GiB the other can still reserve up to 4 GiB.
+- **The pool's total against `pool_size`.** The shares alone do not bound the total, because a
+ consumer keeps what it reserved while fewer consumers were registered. `pool_size` is computed
+ for the executor but applied to each task's pool, so Spark's own limit on the task is at least as
+ tight unless the deprecated `spark.comet.exec.memoryPool.fraction` is below `1.0`.
+
+Two details matter for tuning:
-This is why `fair_unified` spills earlier than `greedy_unified`. It is not a per-consumer quota, and
-reading it as one overstates the memory a multi-operator task can use.
+- The pool is task-shared (see below), so `num_consumers` counts every registered consumer across
+ every native plan in the task, not just the plan making the request. Registering a consumer
+ lowers every other consumer's share, but does not take back memory they already hold.
+- Every registered consumer gets a share, whether or not it can spill. DataFusion's
+ `FairSpillPool` divides the pool among spillable consumers only.
+ [Issue #5465](https://github.com/apache/datafusion-comet/issues/5465) tracks doing the same here.
+
+This is why `fair_unified` can spill earlier than `greedy_unified`: a consumer at its share is
+refused even when the rest of the pool is free, which keeps that memory for the task's other
+consumers.
### Task-shared pools and their lifetime
diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md
index e1257012ec..edfc5bc82a 100644
--- a/docs/source/user-guide/latest/tuning.md
+++ b/docs/source/user-guide/latest/tuning.md
@@ -133,10 +133,17 @@ pre-shuffle operators and one for the shuffle writer). The shared pool ensures t
memory usage stays within the per-task limit.
The `fair_unified` pool prevents operators from using more than an even fraction of the available memory
-(i.e. `pool_size / num_reservations`). This pool works best when you know beforehand
+(i.e. `pool_size / num_consumers`, where `num_consumers` counts the memory consumers registered by all of the task's
+native plans). This pool works best when you know beforehand
the query has multiple operators that will likely all need to spill. Sometimes it will cause spills even
when there is sufficient memory in order to leave enough memory for other operators.
+Comet 0.15.0 through 1.0.0 capped the memory of all of a task's operators combined at one operator's share, because of a
+bug ([#5961](https://github.com/apache/datafusion-comet/issues/5961)). Tasks with several operators can now reserve more
+memory before they spill than they could in those releases. The difference is largest on executors that run few tasks
+at once, where Spark's own limit on each task is loosest. If you sized executor memory against one of those releases,
+check that executors still have enough headroom; see [Sizing the Overhead from the Memory Usage Log].
+
The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not
need to spill or have a single spillable operator.
diff --git a/native/core/src/execution/memory_pools/fair_pool.rs b/native/core/src/execution/memory_pools/fair_pool.rs
index 60e5d51c38..2fbd8224bc 100644
--- a/native/core/src/execution/memory_pools/fair_pool.rs
+++ b/native/core/src/execution/memory_pools/fair_pool.rs
@@ -16,6 +16,7 @@
// under the License.
use std::{
+ collections::HashMap,
fmt::{Debug, Display, Formatter, Result as FmtResult},
sync::Arc,
};
@@ -41,7 +42,21 @@ pub struct CometFairMemoryPool {
struct CometFairPoolState {
used: usize,
- num: usize,
+ /// Bytes held by each registered consumer, keyed by [`MemoryConsumer::id`]. The sibling
+ /// reservations that `new_empty()`, `split()` and `take()` create belong to the same consumer,
+ /// so they count against one share. The pool keeps these totals itself because
+ /// `reservation.size()` covers only one reservation, and DataFusion updates it after calling
+ /// `try_grow` but before calling `shrink`.
+ consumers: HashMap<usize, usize>,
+}
+
+impl CometFairPoolState {
+ /// The bytes held by `reservation`'s consumer across all of its reservations.
+ fn consumer_used(&mut self, reservation: &MemoryReservation) -> &mut usize {
+ self.consumers
+ .get_mut(&reservation.consumer().id())
+ .expect("reservation's consumer is not registered with the pool")
+ }
}
impl Debug for CometFairMemoryPool {
@@ -50,7 +65,7 @@ impl Debug for CometFairMemoryPool {
f.debug_struct("CometFairMemoryPool")
.field("pool_size", &self.pool_size)
.field("used", &state.used)
- .field("num", &state.num)
+ .field("num", &state.consumers.len())
.field("overcommit", &self.spark.overcommit())
.finish()
}
@@ -72,7 +87,10 @@ impl CometFairMemoryPool {
Self {
spark,
pool_size,
- state: Mutex::new(CometFairPoolState { used: 0, num: 0 }),
+ state: Mutex::new(CometFairPoolState {
+ used: 0,
+ consumers: HashMap::new(),
+ }),
}
}
}
@@ -85,7 +103,7 @@ impl Display for CometFairMemoryPool {
"CometFairMemoryPool(pool_size={}, used={}, num={}, overcommit={})",
self.pool_size,
state.used,
- state.num,
+ state.consumers.len(),
self.spark.overcommit()
)
}
@@ -96,74 +114,76 @@ impl MemoryPool for CometFairMemoryPool {
"CometFairMemoryPool"
}
- fn register(&self, _: &MemoryConsumer) {
- let mut state = self.state.lock();
- state.num = state
- .num
- .checked_add(1)
- .expect("unexpected amount of register happened");
+ fn register(&self, consumer: &MemoryConsumer) {
+ self.state.lock().consumers.insert(consumer.id(), 0);
}
- fn unregister(&self, _: &MemoryConsumer) {
- let mut state = self.state.lock();
- state.num = state
- .num
- .checked_sub(1)
- .expect("unexpected amount of unregister happened");
+ fn unregister(&self, consumer: &MemoryConsumer) {
+ // DataFusion unregisters a consumer after its last reservation has dropped and released
+ // its bytes. If a release panicked, this runs while unwinding, so it must not panic too.
+ self.state.lock().consumers.remove(&consumer.id());
}
- /// Records memory that already exists, so it must not fail and ignores the fair limit.
+ /// Records memory that already exists, so it must not fail and ignores the fair and pool
+ /// limits.
/// See [`SparkMemory`].
- fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
+ fn grow(&self, reservation: &MemoryReservation, additional: usize) {
if additional == 0 {
return;
}
let mut state = self.state.lock();
self.spark.acquire(additional);
state.used = state.used.saturating_add(additional);
+ let consumer_used = state.consumer_used(reservation);
+ *consumer_used = consumer_used.saturating_add(additional);
}
- fn shrink(&self, _reservation: &MemoryReservation, subtractive: usize) {
+ fn shrink(&self, reservation: &MemoryReservation, subtractive: usize) {
if subtractive > 0 {
let mut state = self.state.lock();
- // We don't use reservation.size() here because DataFusion 53+ decrements
- // the reservation's atomic size before calling pool.shrink(), so it would
- // reflect the post-shrink value rather than the pre-shrink value.
- if state.used < subtractive {
+ let consumer_used = *state.consumer_used(reservation);
+ if consumer_used < subtractive {
panic!(
- "Failed to release {subtractive} bytes where only {} bytes tracked by pool",
- state.used
+ "Failed to release {subtractive} bytes where only {consumer_used} bytes tracked for the consumer"
)
}
self.spark
.release(subtractive)
.unwrap_or_else(|_| panic!("Failed to release {subtractive} bytes"));
state.used = state.used.checked_sub(subtractive).unwrap();
+ *state.consumer_used(reservation) -= subtractive;
}
}
fn try_grow(
&self,
- _reservation: &MemoryReservation,
+ reservation: &MemoryReservation,
additional: usize,
) -> Result<(), DataFusionError> {
if additional > 0 {
let mut state = self.state.lock();
- let num = state.num;
+ let num = state.consumers.len();
let limit = self
.pool_size
.checked_div(num)
.expect("overflow in checked_div");
- // We use state.used instead of reservation.size() because DataFusion 53+
- // calls pool.try_grow() before incrementing the reservation's atomic size,
- // so reservation.size() would not include prior grows.
- let used = state.used;
- if limit < used + additional {
+ let consumer_used = *state.consumer_used(reservation);
+ if limit < consumer_used.saturating_add(additional) {
return resources_err!(
- "Failed to acquire {additional} bytes where {used} bytes already reserved ({} bytes overcommitted) and the fair limit is {limit} bytes, {num} registered",
+ "Failed to acquire {additional} bytes where this consumer already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num} registered ({} bytes overcommitted)",
self.spark.overcommit()
);
}
+ // The shares alone do not bound the pool's total, because a consumer keeps what it
+ // reserved before another consumer registered.
+ let used = state.used;
+ if self.pool_size < used.saturating_add(additional) {
+ return resources_err!(
+ "Failed to acquire {additional} bytes where {used} bytes already reserved ({} bytes overcommitted) and the pool limit is {} bytes",
+ self.spark.overcommit(),
+ self.pool_size
+ );
+ }
// A partial grant is handed back and refused, which triggers spilling in the caller.
if let Err(refusal) = self.spark.try_acquire(additional)? {
@@ -179,6 +199,7 @@ impl MemoryPool for CometFairMemoryPool {
.used
.checked_add(additional)
.expect("overflow in checked_add");
+ *state.consumer_used(reservation) += additional;
}
Ok(())
}
@@ -211,4 +232,108 @@ mod tests {
// Spark gets back exactly the 100 bytes it granted.
assert_eq!(fake.released(), vec![100]);
}
+
+ #[test]
+ fn each_consumer_is_limited_to_its_own_share() {
+ // Spark grants everything, so only the pool's own checks refuse.
+ let fake = FakeSpark::with(usize::MAX);
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 100));
+ let first = MemoryConsumer::new("first").register(&pool);
+ let second = MemoryConsumer::new("second").register(&pool);
+
+ // Each consumer's share is 50 bytes, whatever the other one holds.
+ first.try_grow(40).unwrap();
+ second.try_grow(20).unwrap();
+ first.try_grow(10).unwrap();
+ assert!(first.try_grow(1).is_err());
+ second.try_grow(30).unwrap();
+ assert!(second.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 100);
+ assert_eq!(fake.held(), 100);
+ }
+
+ #[test]
+ fn a_consumer_registered_late_is_limited_by_the_pool_total() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 90));
+ let first = MemoryConsumer::new("first").register(&pool);
+ // Alone, the first consumer's share is the whole pool.
+ first.try_grow(60).unwrap();
+
+ // A second consumer halves both shares, but the first keeps the 60 bytes it holds.
+ let second = MemoryConsumer::new("second").register(&pool);
+ assert!(first.try_grow(1).is_err());
+ second.try_grow(30).unwrap();
+ // The second consumer is 15 bytes under its share, but the pool is full.
+ assert!(second.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 90);
+
+ // Once the first consumer releases memory, the second can use the rest of its share.
+ first.shrink(30);
+ second.try_grow(15).unwrap();
+ assert!(second.try_grow(1).is_err());
+
+ // Unregistering the second consumer gives the first the whole pool again.
+ drop(second);
+ first.try_grow(60).unwrap();
+ assert_eq!(pool.reserved(), 90);
+ assert_eq!(fake.held(), 90);
+ }
+
+ #[test]
+ fn sibling_reservations_draw_on_one_share() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 100));
+ let first = MemoryConsumer::new("first").register(&pool);
+ // Like the reservation that a sort's streaming merge creates for each batch it reads.
+ let sibling = first.new_empty();
+ let second = MemoryConsumer::new("second").register(&pool);
+
+ // Both of the first consumer's reservations draw on its one 50-byte share.
+ first.try_grow(40).unwrap();
+ assert!(sibling.try_grow(40).is_err());
+ sibling.try_grow(10).unwrap();
+ assert!(first.try_grow(1).is_err());
+ assert!(sibling.try_grow(1).is_err());
+
+ // So the second consumer can still reserve its whole share.
+ second.try_grow(50).unwrap();
+ assert_eq!(pool.reserved(), 100);
+
+ // Dropping a reservation returns its bytes to the consumer's share. The consumer stays
+ // registered while its other reservation lives.
+ drop(first);
+ sibling.try_grow(40).unwrap();
+ assert!(sibling.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 100);
+ assert_eq!(fake.held(), 100);
+ }
+
+ #[test]
+ fn split_and_take_keep_the_bytes_on_their_consumer() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 100));
+ let mut first = MemoryConsumer::new("first").register(&pool);
+ let _second = MemoryConsumer::new("second").register(&pool);
+
+ first.try_grow(50).unwrap();
+ let split = first.split(20);
+ let taken = first.take();
+ // The consumer's 50 bytes now sit in split and taken, and the emptied reservation gets no
+ // share of its own.
+ for reservation in [&first, &split, &taken] {
+ assert!(reservation.try_grow(1).is_err());
+ }
+
+ // Shrinking one reservation makes room in the share for another.
+ split.shrink(10);
+ first.try_grow(10).unwrap();
+ assert!(taken.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 50);
+ assert_eq!(fake.held(), 50);
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org