From 8bcacb60688e0b3d03e7551fd24cd2eca7f2d56a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 04:38:49 +0000 Subject: [PATCH] fix: let a spilled final aggregate read its spill files back past its memory share (#6544) * fix: let a spilled final aggregate read its spill files back past its memory share DataFusion 55's FinalHashAggregateStream replays its merged spill files through a stream that can't spill, so a refused memory request there fails the task. The merge's read buffers are a sibling reservation of the same consumer and take as many spill files as fit, so the replay often finds the consumer's share already taken (#6254). Both Comet pools now record that request as overcommit instead of refusing it. They recognize it as a request from a FinalHashAggregateStream consumer while another of its reservations holds memory, which in DataFusion 55.1 happens only during the replay. The replay emits its finished groups after every batch, so the overcommit stays around one batch of groups, and releases repay it first. Refusals while the aggregate reads its input, and while the merge picks its files, are unchanged. Remove this once Comet's DataFusion includes apache/datafusion#25383. * fix: let an ordered final aggregate read its spill files back too DataFusion runs a final aggregate whose input is sorted on some of its grouping keys as an OrderedFinalAggregateStream. It merges and replays its spill files the same way FinalHashAggregateStream does, so its replay failed the task the same way. Treat its consumer as a final aggregate too. Explain in spill_replay.rs why recording the replay's request is safe: the replay asks for memory only after it has aggregated a batch, so the memory already exists, as it does for a grow. Describe the exception in the memory management guide, which said the fair pool always refuses a request that fails its local checks. * docs: tighten the spill replay comments and guide The fair pool is the only one with a share and a pool total to skip, so say so. The greedy pool takes its tracking lock for other consumers when Spark refuses them, not never. Point whoever removes the workaround at the tests that show whether the replay still needs it. * test: share the spill replay tests between the pools fair_pool.rs and unified_pool.rs each had a copy of the same two spill replay scenarios. Run them once, in spill_replay.rs, against both pool types as createPlan builds them, which also checks that the wrappers pass register through to the greedy pool's tracking. Keep a greedy pool test for its per-consumer map, and drop a test import that the pool's own imports now cover. Share the final aggregate spill checks between the two #6254 tests in CometAggregateSuite, and use checkCometAnswer, which collects once and labels the Comet answer as Comet's. * refactor: move the spill replay workaround into a pool wrapper The workaround for #6254 lived in both Comet pools. The fair pool sent its three refusals through a check, and the greedy pool tracked each final aggregate across its reservations. #5613 reworks the fair pool's try_grow, and #6583 has to take the workaround out again, so both would have had to rework the pools. SpillReplayPool, in spill_replay.rs, now wraps the tracked Comet pool instead. It keeps a total for each final aggregate's consumer. When the pool refuses a request from one while another of its reservations holds memory, it calls the pool's grow, which skips the fair pool's limits and carries what Spark doesn't grant as overcommit. fair_pool.rs and unified_pool.rs are back to main's versions, and overcommit() looks through the wrapper. The Rust tests run against both pool types as createPlan builds them. They now reach the fair pool's pool-limit refusal too, and check that a failed Spark call during the replay is returned rather than recorded. The guide adds the wrapper to the pool stack and keeps a short section. The DataFusion 55.1 facts that an upgrade has to re-check stay in spill_replay.rs, which names #6583 next to the removal condition. (cherry picked from commit ba7c89203813b848fb620f1078276b9f42614f96) Adapted for branch-1.1: - The tests build the pools over a fake Spark through create_pool, a seam that #6271 added to memory_pools. Only that seam is ported, not #6271's memory usage log change, which is not on branch-1.1: create_memory_pool builds the pools through create_pool, with_spark is pub(super), and the pools' new constructors, which nothing calls any more, are removed. - overcommit() and unwrap_task_shared(), which only that log reads, are not ported, so nothing looks through the wrapper. spill_replay.rs drops unwrap_spill_replay, and its tests drop their overcommit(&pool) assertions, which restate pool.reserved() minus what the fake Spark holds. The test that shrinks the replay checks pool.reserved() instead. - memory_management.md: the paragraph about the task's shared CometTaskMemoryManager, from #6261, is not on branch-1.1, so the new section follows the task-shared pools section directly. - CometAggregateSuite: only its imports differ. This branch does not import Column, CometBaseAggregateExec or CometSortAggregateExec. --- .../contributor-guide/memory_management.md | 25 +- .../src/execution/memory_pools/fair_pool.rs | 17 +- native/core/src/execution/memory_pools/mod.rs | 24 +- .../execution/memory_pools/spill_replay.rs | 360 ++++++++++++++++++ .../execution/memory_pools/unified_pool.rs | 19 +- .../comet/exec/CometAggregateSuite.scala | 92 ++++- 6 files changed, 494 insertions(+), 43 deletions(-) create mode 100644 native/core/src/execution/memory_pools/spill_replay.rs diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 9768698b0d8..2ab80ed773b 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -243,14 +243,17 @@ ignored and the pool is always `UnboundedMemoryPool`. the inside out, a Comet plan in the default configuration sees: ```text -[LoggingMemoryPool] <- only when spark.comet.debug.memory=true - [TaskSharedMemoryPool] <- RAII handle for the per-task registry - [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages - [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI +[LoggingMemoryPool] <- only when spark.comet.debug.memory=true + [TaskSharedMemoryPool] <- RAII handle for the per-task registry + [SpillReplayPool] <- lets a spilled final aggregate read its spill files back + [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages + [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI ``` Each decorator forwards every `MemoryPool` method to its inner pool, so `reserved()` at any level -reports the base pool's number. +reports the base pool's number. `SpillReplayPool` also turns one kind of refused `try_grow` into a +`grow`; see +[Final aggregates reading their spill files back](#final-aggregates-reading-their-spill-files-back). ### The unified pools @@ -330,6 +333,18 @@ forget, and a `createPlan` that fails partway through cleans up on unwind. insert a replacement before the dying pool reaches the registry lock. The drop therefore compares pointers and only removes an entry that is still its own. +### Final aggregates reading their spill files back + +`SpillReplayPool` (`spill_replay.rs`) wraps both Comet pools to work around +[issue #6254](https://github.com/apache/datafusion-comet/issues/6254). Once one of DataFusion 55's +final aggregates has spilled, it reads its spill files back through an aggregate that cannot spill, +so a refused `try_grow` there fails the task. When the pool refuses such a request, +`SpillReplayPool` records it with the pool's `grow` instead, which skips `CometFairMemoryPool`'s +local checks and carries what Spark does not grant as overcommit. Every other refusal is passed on +unchanged. `spill_replay.rs` describes how the wrapper recognizes these requests and what that relies +on in DataFusion. [Issue #6583](https://github.com/apache/datafusion-comet/issues/6583) tracks +removing it. + ## How DataFusion consumes the pool Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservation` API: diff --git a/native/core/src/execution/memory_pools/fair_pool.rs b/native/core/src/execution/memory_pools/fair_pool.rs index 2fbd8224bc3..c36d2281b16 100644 --- a/native/core/src/execution/memory_pools/fair_pool.rs +++ b/native/core/src/execution/memory_pools/fair_pool.rs @@ -18,11 +18,8 @@ use std::{ collections::HashMap, fmt::{Debug, Display, Formatter, Result as FmtResult}, - sync::Arc, }; -use jni::objects::{Global, JObject}; - use super::spark_memory::SparkMemory; use datafusion::common::resources_err; use datafusion::execution::memory_pool::MemoryConsumer; @@ -72,18 +69,7 @@ impl Debug for CometFairMemoryPool { } impl CometFairMemoryPool { - pub fn new( - task_memory_manager_handle: Arc>>, - pool_size: usize, - task_attempt_id: i64, - ) -> CometFairMemoryPool { - Self::with_spark( - SparkMemory::new(task_memory_manager_handle, task_attempt_id), - pool_size, - ) - } - - fn with_spark(spark: SparkMemory, pool_size: usize) -> CometFairMemoryPool { + pub(super) fn with_spark(spark: SparkMemory, pool_size: usize) -> CometFairMemoryPool { Self { spark, pool_size, @@ -213,6 +199,7 @@ impl MemoryPool for CometFairMemoryPool { mod tests { use super::super::spark_memory::fake::FakeSpark; use super::*; + use std::sync::Arc; #[test] fn grow_past_the_fair_limit_is_recorded_and_refuses_the_next_try_grow() { diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 69311698b2e..373313fd882 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -19,12 +19,15 @@ mod config; mod fair_pool; pub mod logging_pool; mod spark_memory; +mod spill_replay; mod task_shared; mod unified_pool; use datafusion::execution::memory_pool::{MemoryPool, TrackConsumersPool, UnboundedMemoryPool}; use fair_pool::CometFairMemoryPool; use jni::objects::{Global, JObject}; +use spark_memory::SparkMemory; +use spill_replay::SpillReplayPool; use std::num::NonZeroUsize; use std::sync::Arc; use unified_pool::CometUnifiedMemoryPool; @@ -40,6 +43,18 @@ pub(crate) fn create_memory_pool( memory_pool_config: &MemoryPoolConfig, comet_task_memory_manager: Arc>>, task_attempt_id: i64, +) -> Arc { + create_pool(memory_pool_config, task_attempt_id, || { + SparkMemory::new(comet_task_memory_manager, task_attempt_id) + }) +} + +/// Creates the pool that [`create_memory_pool`] does, with `spark` connecting it to Spark's memory +/// manager, so that tests can connect it to a fake instead. +fn create_pool( + memory_pool_config: &MemoryPoolConfig, + task_attempt_id: i64, + spark: impl FnOnce() -> SparkMemory, ) -> Arc { const NUM_TRACKED_CONSUMERS: usize = 10; @@ -55,16 +70,15 @@ pub(crate) fn create_memory_pool( match pool_type { MemoryPoolType::GreedyUnified => acquire_task_shared_pool(task_attempt_id, || { - tracked(CometUnifiedMemoryPool::new( - comet_task_memory_manager, + Arc::new(SpillReplayPool::new( task_attempt_id, + tracked(CometUnifiedMemoryPool::with_spark(spark())), )) }), MemoryPoolType::FairUnified => acquire_task_shared_pool(task_attempt_id, || { - tracked(CometFairMemoryPool::new( - comet_task_memory_manager, - pool_size, + Arc::new(SpillReplayPool::new( task_attempt_id, + tracked(CometFairMemoryPool::with_spark(spark(), pool_size)), )) }), MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()), diff --git a/native/core/src/execution/memory_pools/spill_replay.rs b/native/core/src/execution/memory_pools/spill_replay.rs new file mode 100644 index 00000000000..069b61ccb83 --- /dev/null +++ b/native/core/src/execution/memory_pools/spill_replay.rs @@ -0,0 +1,360 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Lets a final aggregate read its spill files back past its share of memory (#6254). +//! +//! Once one of DataFusion 55's final aggregates has spilled, it merges its sorted spill files and +//! replays them through an `OrderedFinalAggregateStream` that has no way to spill, so a refused +//! memory request there fails the task. `FinalHashAggregateStream` does this, and so does +//! `OrderedFinalAggregateStream` itself, which DataFusion uses when the input is sorted on some of +//! the grouping keys. The merge reserves read buffers for as many spill files as fit, and those +//! buffers belong to the same consumer, so the replay often finds the consumer's share already +//! taken. +//! +//! [`SpillReplayPool`] wraps a Comet pool. When the pool refuses a request from the replay, it +//! records the request with the pool's `grow` instead, which ignores `CometFairMemoryPool`'s limits +//! and carries what Spark doesn't grant as overcommit. +//! +//! This relies on the following in DataFusion 55.1, which a DataFusion upgrade has to re-check: +//! +//! - The two aggregates name their consumers `FinalHashAggregateStream[{partition}]` and +//! `OrderedFinalAggregateStream[{partition}]`, and their merge and replay reserve through sibling +//! reservations of that consumer. +//! - A final aggregate grows one reservation while another of its reservations holds memory only +//! during the replay, while the merge holds its read buffers. Before that, the aggregate's table +//! is its only reservation holding memory, so a refusal stands and makes it spill. The merge +//! picks its files while nothing else is held, so a refusal there still limits how many it opens. +//! - The replay asks for memory only after it has aggregated a batch, so like a `grow`, the request +//! is for memory that already exists. +//! - The replay emits every finished group after each batch, so it holds about one batch of groups, +//! and releasing memory repays the overcommit first. +//! +//! Remove this, as #6583 describes, once Comet's DataFusion has apache/datafusion#25383, which +//! leaves the replay room when the merge picks its files. The #6254 tests in `CometAggregateSuite` +//! fail without this on DataFusion 55.1, so they show whether the replay still needs it. + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use log::debug; +use parking_lot::Mutex; + +/// Wraps a Comet pool and records the spill replay requests that it refuses; see the +/// [module documentation](self). Every other call, and every other refusal, is passed through +/// unchanged. +#[derive(Debug)] +pub(super) struct SpillReplayPool { + task_attempt_id: i64, + inner: Arc, + /// Bytes held by each final aggregate's consumer across all of its reservations, keyed by + /// [`MemoryConsumer::id`], because `reservation.size()` covers only one reservation. Other + /// consumers aren't tracked, so they never take the lock. + final_aggregates: Mutex>, +} + +impl SpillReplayPool { + pub(super) fn new(task_attempt_id: i64, inner: Arc) -> Self { + Self { + task_attempt_id, + inner, + final_aggregates: Mutex::new(HashMap::new()), + } + } + + /// Applies `update` to what `reservation`'s consumer holds, if it is a final aggregate. + fn track(&self, reservation: &MemoryReservation, update: impl FnOnce(&mut usize)) { + if is_final_aggregate(reservation.consumer()) { + if let Some(used) = self + .final_aggregates + .lock() + .get_mut(&reservation.consumer().id()) + { + update(used); + } + } + } + + /// Whether a refused request from `reservation` comes from a final aggregate reading its spill + /// files back, the only time one of its reservations grows while another holds memory. + fn is_spill_replay(&self, reservation: &MemoryReservation) -> bool { + is_final_aggregate(reservation.consumer()) + && self + .final_aggregates + .lock() + .get(&reservation.consumer().id()) + .is_some_and(|&consumer_used| consumer_used > reservation.size()) + } +} + +impl fmt::Display for SpillReplayPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self.inner.as_ref(), f) + } +} + +impl MemoryPool for SpillReplayPool { + fn name(&self) -> &str { + self.inner.name() + } + + fn register(&self, consumer: &MemoryConsumer) { + if is_final_aggregate(consumer) { + self.final_aggregates.lock().insert(consumer.id(), 0); + } + self.inner.register(consumer) + } + + fn unregister(&self, consumer: &MemoryConsumer) { + if is_final_aggregate(consumer) { + self.final_aggregates.lock().remove(&consumer.id()); + } + self.inner.unregister(consumer) + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.track(reservation, |used| *used = used.saturating_add(additional)); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.track(reservation, |used| *used = used.saturating_sub(shrink)); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + match self.inner.try_grow(reservation, additional) { + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(refusal)) + if self.is_spill_replay(reservation) => + { + debug!( + "Task {} records {additional} bytes for {} while it reads its spill files back: {refusal}", + self.task_attempt_id, + reservation.consumer().name() + ); + self.inner.grow(reservation, additional); + } + Err(e) => return Err(e), + } + self.track(reservation, |used| *used = used.saturating_add(additional)); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +/// Whether `consumer` belongs to a final aggregate whose spill replay can't spill. +fn is_final_aggregate(consumer: &MemoryConsumer) -> bool { + let name = consumer.name(); + name.starts_with("FinalHashAggregateStream[") + || name.starts_with("OrderedFinalAggregateStream[") +} + +#[cfg(test)] +mod tests { + use super::super::spark_memory::fake::FakeSpark; + use super::super::{create_pool, MemoryPoolConfig, MemoryPoolType}; + use super::*; + use datafusion::execution::memory_pool::UnboundedMemoryPool; + + /// A pool of each type, built the way `createPlan` builds it and connected to a fake Spark + /// that grants at most 100 bytes. The fair pool's own limits are far above that, so only + /// Spark refuses. Task-shared pools are keyed by task attempt process-wide, so each pool needs + /// its own id. + fn each_pool_type( + task_attempt_ids: [i64; 2], + ) -> Vec<(&'static str, Arc, Arc)> { + [ + ("greedy_unified", MemoryPoolType::GreedyUnified), + ("fair_unified", MemoryPoolType::FairUnified), + ] + .into_iter() + .zip(task_attempt_ids) + .map(|((name, pool_type), task_attempt_id)| { + let fake = FakeSpark::with(100); + let config = MemoryPoolConfig::new(pool_type, 1000); + let pool = create_pool(&config, task_attempt_id, || fake.memory()); + (name, pool, fake) + }) + .collect() + } + + /// A `fair_unified` pool of `pool_size` bytes, built the way `createPlan` builds it and + /// connected to a fake Spark that grants everything, so only the pool's own limits refuse. + fn fair_pool(task_attempt_id: i64, pool_size: usize) -> (Arc, Arc) { + let fake = FakeSpark::with(usize::MAX); + let config = MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size); + let pool = create_pool(&config, task_attempt_id, || fake.memory()); + (pool, fake) + } + + #[test] + fn a_final_aggregate_reading_its_spill_files_back_carries_what_spark_refuses() { + for (name, pool, fake) in each_pool_type([-3011, -3012]) { + // Like DataFusion 55's final hash aggregate, whose spill merge and replay table are + // sibling reservations of one consumer. + let merge = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&pool); + let replay = merge.new_empty(); + merge.try_grow(90).unwrap(); + + // Spark grants 10 of the replay's 30 bytes, and the other 20 are overcommit. + replay.try_grow(30).unwrap(); + assert_eq!(pool.reserved(), 120, "{name}"); + assert_eq!(fake.held(), 100, "{name}"); + + // The replay emits groups and shrinks, which repays the overcommit before Spark. + replay.shrink(25); + assert_eq!(pool.reserved(), 95, "{name}"); + assert_eq!(fake.held(), 95, "{name}"); + drop(merge); + drop(replay); + assert_eq!(pool.reserved(), 0, "{name}"); + assert_eq!(fake.held(), 0, "{name}"); + } + } + + #[test] + fn a_final_aggregate_reading_its_spill_files_back_may_pass_its_fair_limit() { + let (pool, fake) = fair_pool(-3013, 100); + let merge = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&pool); + let replay = merge.new_empty(); + merge.try_grow(90).unwrap(); + + // The aggregate is the only consumer, so its share is the whole pool. + replay.try_grow(30).unwrap(); + assert_eq!(pool.reserved(), 120); + assert_eq!(fake.held(), 120); + } + + #[test] + fn a_final_aggregate_reading_its_spill_files_back_may_pass_the_pool_limit() { + let (pool, fake) = fair_pool(-3014, 100); + let sort = MemoryConsumer::new("ExternalSorter[0]").register(&pool); + sort.try_grow(70).unwrap(); + let merge = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&pool); + let replay = merge.new_empty(); + merge.try_grow(25).unwrap(); + + // 25 + 20 is within the aggregate's share of 50, but 70 + 25 + 20 is over the pool's 100. + replay.try_grow(20).unwrap(); + assert_eq!(pool.reserved(), 115); + assert_eq!(fake.held(), 115); + } + + #[test] + fn other_refusals_are_unchanged() { + for (name, pool, fake) in each_pool_type([-3015, -3016]) { + // While the aggregate reads its input, its table is the consumer's only reservation + // holding memory, so a refusal makes it spill. + let table = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&pool); + let replay = table.new_empty(); + table.try_grow(90).unwrap(); + assert!(table.try_grow(30).is_err(), "{name}"); + drop(table); + + // The merge picks its spill files while nothing else is held, so a refusal still + // limits how many it opens. + let merge = replay.new_empty(); + merge.try_grow(90).unwrap(); + assert!(merge.try_grow(30).is_err(), "{name}"); + drop(merge); + + // Any other operator with a sibling holding memory is still refused. + let sort = MemoryConsumer::new("ExternalSorterMerge[0]").register(&pool); + let sibling = sort.new_empty(); + sort.try_grow(90).unwrap(); + assert!(sibling.try_grow(30).is_err(), "{name}"); + assert_eq!(pool.reserved(), 90, "{name}"); + assert_eq!(fake.held(), 90, "{name}"); + } + } + + #[test] + fn a_failed_spark_call_during_the_replay_is_not_recorded() { + for (name, pool_type, task_attempt_id) in [ + ("greedy_unified", MemoryPoolType::GreedyUnified, -3017), + ("fair_unified", MemoryPoolType::FairUnified, -3018), + ] { + let fake = FakeSpark::failing(); + let config = MemoryPoolConfig::new(pool_type, 1000); + let pool = create_pool(&config, task_attempt_id, || fake.memory()); + let merge = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&pool); + let replay = merge.new_empty(); + merge.grow(90); + + // Only a refusal is recorded. A JNI error goes back to the aggregate. + let err = replay.try_grow(30).unwrap_err(); + assert!( + !matches!(err, DataFusionError::ResourcesExhausted(_)), + "{name}: {err}" + ); + assert_eq!(pool.reserved(), 90, "{name}"); + } + } + + #[test] + fn a_final_aggregate_is_tracked_across_its_reservations_until_it_unregisters() { + let pool = Arc::new(SpillReplayPool::new( + -3019, + Arc::new(UnboundedMemoryPool::default()), + )); + let dyn_pool: Arc = Arc::clone(&pool) as _; + let merge = MemoryConsumer::new("FinalHashAggregateStream[0]").register(&dyn_pool); + let replay = merge.new_empty(); + merge.try_grow(90).unwrap(); + replay.grow(20); + replay.shrink(5); + let sort = MemoryConsumer::new("ExternalSorter[0]").register(&dyn_pool); + sort.grow(10); + let id = merge.consumer().id(); + assert_eq!(*pool.final_aggregates.lock(), HashMap::from([(id, 105)])); + + drop(merge); + drop(replay); + assert!(pool.final_aggregates.lock().is_empty()); + } + + #[test] + fn only_final_aggregates_replay_their_spill_files() { + for name in [ + "FinalHashAggregateStream[3]", + "OrderedFinalAggregateStream[0]", + ] { + assert!(is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); + } + // Partial aggregates emit early instead of spilling, and Comet never plans single ones. + for name in [ + "PartialHashAggregateStream[0]", + "OrderedPartialAggregateStream[0]", + "SingleHashAggregateStream[0]", + "ExternalSorterMerge[0]", + ] { + assert!(!is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); + } + } +} diff --git a/native/core/src/execution/memory_pools/unified_pool.rs b/native/core/src/execution/memory_pools/unified_pool.rs index a9f6a0b4227..07b84659b5f 100644 --- a/native/core/src/execution/memory_pools/unified_pool.rs +++ b/native/core/src/execution/memory_pools/unified_pool.rs @@ -17,10 +17,7 @@ use std::{ fmt::{Debug, Display, Formatter, Result as FmtResult}, - sync::{ - atomic::{AtomicUsize, Ordering::Relaxed}, - Arc, - }, + sync::atomic::{AtomicUsize, Ordering::Relaxed}, }; use super::spark_memory::SparkMemory; @@ -28,7 +25,6 @@ use datafusion::{ common::{resources_datafusion_err, DataFusionError}, execution::memory_pool::{MemoryPool, MemoryReservation}, }; -use jni::objects::{Global, JObject}; use log::warn; /// A DataFusion `MemoryPool` implementation for Comet that delegates to @@ -49,17 +45,7 @@ impl Debug for CometUnifiedMemoryPool { } impl CometUnifiedMemoryPool { - pub fn new( - task_memory_manager_handle: Arc>>, - task_attempt_id: i64, - ) -> CometUnifiedMemoryPool { - Self::with_spark(SparkMemory::new( - task_memory_manager_handle, - task_attempt_id, - )) - } - - fn with_spark(spark: SparkMemory) -> CometUnifiedMemoryPool { + pub(super) fn with_spark(spark: SparkMemory) -> CometUnifiedMemoryPool { Self { spark, used: AtomicUsize::new(0), @@ -171,6 +157,7 @@ mod tests { use super::super::spark_memory::fake::FakeSpark; use super::*; use datafusion::execution::memory_pool::MemoryConsumer; + use std::sync::Arc; #[test] fn grow_past_spark_is_recorded_and_refuses_try_grow_until_repaid() { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 92536d154ad..0885ddc49c9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -26,12 +26,12 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.{CometListenerBusUtils, SparkConf} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} -import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.{CometTestBase, DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} -import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec} +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} @@ -154,6 +154,94 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + /** + * Returns the native final aggregates in the plan that `df` last ran, after checking that there + * is one and that each spilled. + */ + private def spilledFinalAggregates(df: DataFrame): Seq[CometHashAggregateExec] = { + val plan = df.queryExecution.executedPlan + val finalAggregates = collect(plan) { + case aggregate: CometHashAggregateExec if aggregate.modes.contains(Final) => aggregate + } + assert(finalAggregates.nonEmpty, s"Expected a native final aggregate:\n$plan") + assert( + finalAggregates.forall(_.metrics("spill_count").value > 0), + "The final aggregate did not spill") + finalAggregates + } + + test("final aggregate that has spilled reads its spill files back (issue #6254)") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 125000, 1, 1) + .selectExpr("lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v") + .write + .parquet(path) + // A 3 MiB pool (3/2048 of the suite's 2g off-heap) makes the final aggregate spill several + // runs of its 62,500 groups. Merging the runs then takes nearly all of its share, and + // reading them back must not fail the task. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (3.0 / 2048).toString) { + val df = spark.read.parquet(path).groupBy("k").agg(expr("count(1)"), expr("max(v)")) + val expected = (0 until 62500).map { i => + val key = i.toString + Row("k" * (128 - key.length) + key, 2L, i + 62500L) + } + // checkCometAnswer runs the query once, so the plan read below is the one that ran. + checkCometAnswer(df, expected) + spilledFinalAggregates(df) + } + } + } + + test("ordered final aggregate that has spilled reads its spill files back (issue #6254)") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 125000, 1, 1) + .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v") + .write + .parquet(path) + // Sorting on g, the first grouping key, makes DataFusion run the final aggregate as an + // OrderedFinalAggregateStream, which spills and reads its spill files back the same way. + // Spark drops a sort under an aggregate unless EliminateSorts is excluded. With a 2.5 MiB + // pool, merging the runs takes nearly all of the aggregate's share. The sort's merge + // reservation is lowered so that the sort fits in the pool too. + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName, + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (2.5 / 2048).toString) { + val df = spark.read + .parquet(path) + .repartition(1, col("g")) + .sortWithinPartitions("g") + .groupBy("g", "k") + .agg(expr("count(1)"), expr("max(v)")) + val expected = (0 until 62500).map { i => + val key = i.toString + Row(0L, "k" * (128 - key.length) + key, 2L, i + 62500L) + } + // checkCometAnswer would compare the rows in order, because the plan has a sort. + // Collecting runs the query once, so the plan read below is the one that ran. + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + + // Only a sort in the same native plan keeps the final aggregate's input sorted. + assert( + spilledFinalAggregates(df).exists(_.child match { + case partial: CometHashAggregateExec => partial.child.isInstanceOf[CometSortExec] + case _ => false + }), + "Expected a native final aggregate over a partial aggregate over a sort:\n" + + df.queryExecution.executedPlan) + } + } + } + test("collect_list over struct with non-nullable fields") { // Building a struct from non-nullable columns yields non-nullable struct fields. Native // collect_list derives its declared output element type from the argument's declared type,