comphead commented on code in PR #7096:
URL: https://github.com/apache/datafusion-comet/pull/7096#discussion_r4338820622
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -203,15 +282,213 @@ impl MemoryPool for CometFairMemoryPool {
#[cfg(test)]
mod tests {
use super::super::spark_memory::fake::FakeSpark;
+ use super::super::task_shared::RunningTask;
use super::*;
use std::sync::Arc;
+ /// A pool of `pool_size` bytes, all of the off-heap pool as with the
default fraction, which
+ /// `tasks` share, connected to `fake`.
+ fn shared_pool(
+ fake: &Arc<FakeSpark>,
+ pool_size: usize,
+ tasks: &RunningTasks,
+ ) -> Arc<dyn MemoryPool> {
+ Arc::new(CometFairMemoryPool::with_spark(
+ fake.memory(),
+ pool_size,
+ pool_size,
+ tasks.clone(),
+ ))
+ }
+
+ /// A pool of `pool_size` bytes for a task that runs alone, connected to
`fake`.
+ fn pool(fake: &Arc<FakeSpark>, pool_size: usize) -> Arc<dyn MemoryPool> {
+ shared_pool(fake, pool_size, &RunningTasks::default())
+ }
+
+ /// Starts `n` tasks, which run until the returned guards drop.
+ fn start(tasks: &RunningTasks, n: usize) -> Vec<RunningTask> {
+ (0..n).map(|_| tasks.start()).collect()
+ }
+
+ /// A consumer that can spill, like a sort's buffer.
+ fn spillable(name: &str, pool: &Arc<dyn MemoryPool>) -> MemoryReservation {
+ MemoryConsumer::new(name)
+ .with_can_spill(true)
+ .register(pool)
+ }
+
+ #[test]
+ fn consumers_that_can_spill_share_their_tasks_share_of_the_pool() {
+ // Eight running tasks share a 1040-byte pool, and Spark grants this
one its 130 bytes.
+ let fake = FakeSpark::with(130);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 8);
+ let pool = shared_pool(&fake, 1040, &tasks);
+ let sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ let other_merge =
MemoryConsumer::new("ExternalSorterMerge[1]").register(&pool);
+
+ // The two sorts share the task's 130 bytes, so the first is refused,
and spills, at 65.
+ sort.try_grow(65).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ // That leaves the rest of what Spark grants the task for the other
sort's merge, which
+ // cannot spill.
+ other_merge.try_grow(65).unwrap();
+ assert_eq!(fake.held(), 130);
+ }
+
+ #[test]
+ fn a_consumer_that_cannot_spill_is_limited_only_by_the_pool() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ let _sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[0]").register(&pool);
+
+ // Neither a share of the task's 100 bytes nor the task's share itself
limits the merge.
+ merge.try_grow(400).unwrap();
+ // The pool does, but only binds before Spark does when the deprecated
+ // `spark.comet.exec.memoryPool.fraction` is below 1.0.
+ assert!(merge.try_grow(1).is_err());
+ assert_eq!(fake.held(), 400);
+ }
+
+ #[test]
+ fn consumers_that_can_spill_share_what_the_others_leave() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool = pool(&fake, 100);
+ let first = spillable("ExternalSorter[0]", &pool);
+ let second = spillable("ExternalSorter[1]", &pool);
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[0]").register(&pool);
+ merge.try_grow(40).unwrap();
+
+ // The two sorts share the 60 bytes that the merge leaves.
+ first.try_grow(30).unwrap();
+ assert!(first.try_grow(1).is_err());
+
+ // Their shares grow as the merge releases memory.
+ merge.shrink(40);
+ first.try_grow(20).unwrap();
+ assert!(first.try_grow(1).is_err());
+ second.try_grow(50).unwrap();
+ assert!(second.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 100);
+ }
+
+ #[test]
+ fn
memory_recorded_for_a_consumer_that_cannot_spill_comes_out_of_the_shares() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool = pool(&fake, 100);
+ let sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ // Like a batch read back from a spill file, which already exists.
+ let reader = MemoryConsumer::new("spill reader").register(&pool);
+ reader.grow(40);
+
+ // The two sorts share the 60 bytes left.
+ sort.try_grow(30).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn a_tasks_share_follows_the_number_of_running_tasks() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let mut running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ let sort = spillable("ExternalSorter[0]", &pool);
+
+ // Four running tasks share the 400-byte pool.
+ sort.try_grow(100).unwrap();
+ assert!(sort.try_grow(1).is_err());
+
+ // Once two of them finish, the other two share it.
+ running.truncate(2);
+ sort.try_grow(100).unwrap();
+ assert!(sort.try_grow(1).is_err());
+
+ // A task that starts takes part of the share back, so the sort spills
before it grows.
+ running.push(tasks.start());
+ assert!(sort.try_grow(1).is_err());
+ sort.shrink(150);
+ sort.try_grow(83).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn the_tasks_total_is_checked_against_the_tasks_share_not_the_pool() {
+ // One task has the whole 200-byte pool, and its sort takes half.
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = tasks.start();
+ let pool = shared_pool(&fake, 200, &tasks);
+ let first = spillable("ExternalSorter[0]", &pool);
+ first.try_grow(100).unwrap();
+
+ // Another task starts, which halves the task's share to the 100 bytes
that the sort holds.
+ // A second sort is under its own limit of 50 bytes, but the task's
total is at its share.
+ let _other_task = tasks.start();
+ let second = spillable("ExternalSorter[1]", &pool);
+ let error = second.try_grow(50).unwrap_err().to_string();
+ assert!(error.contains("the task's share is 100 bytes"), "{error}");
+
+ // A consumer that cannot spill is limited by the pool, not by the
task's share.
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[1]").register(&pool);
+ merge.try_grow(100).unwrap();
+ assert_eq!(pool.reserved(), 200);
+ }
+
+ /// The trade-off that the migration guide describes: a consumer that
cannot spill can take
+ /// what Spark grants, which can be the task's whole share, and then a
consumer that can spill
+ /// is refused its first request, so a sort fails where it used to run.
+ #[test]
+ fn
a_consumer_that_can_spill_is_refused_once_the_others_hold_the_tasks_share() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ // Like a shuffled hash join's build side, which Spark lets grow past
`off_heap_size / 4`
+ // while the other tasks hold nothing yet.
+ let build = MemoryConsumer::new("HashJoinInput[0]").register(&pool);
+ build.try_grow(100).unwrap();
+ let sort = spillable("ExternalSorter[0]", &pool);
+
+ let error = sort.try_grow(1).unwrap_err().to_string();
+ assert!(error.contains("the fair limit is 0 bytes"), "{error}");
+
+ // The sort can reserve again once the build side gives memory back.
+ build.shrink(60);
+ sort.try_grow(60).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn a_tasks_share_is_at_most_the_pool() {
Review Comment:
This test looks like a copy of
`the_fraction_limits_a_task_that_shares_the_off_heap_pool` in `mod.rs`:
100-byte pool, 1000-byte off-heap size, two running tasks, and the same
`try_grow(100)` then refused `try_grow(1)` on a sort. The `mod.rs` version goes
through `create_pool`, so it also checks the argument order and the consumer
that cannot spill. Could this one go?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
