Skip to content

Commit 0ee7a4d

Browse files
committed
fix: keep intermediate spill merges within their admitted memory
The amd64 SQL CI failure in aggregate_memory_spill Case G occurred when four final partitions shared a 1 MiB GreedyMemoryPool: replay requested 6.6 KiB with only 948 bytes free. Capping test fan-in removed default-path coverage without protecting competing replay allocations. Trade two-batch read-ahead for more intermediate inputs within the original admitted grant, restore the default fan-in test, and cover the peer-growth regression directly. Keep the original selection and read-ahead for a no-readmission retry and preserve the first error if that retry fails. Flush intermediate output according to retained-source and prospective output memory instead of every input boundary. Cover full batches, short and skewed inputs, dictionaries, ties, and seeded concurrent aggregations. Reuse test fixtures and share spill admission cost calculations.
1 parent 6987df2 commit 0ee7a4d

7 files changed

Lines changed: 856 additions & 316 deletions

File tree

‎datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs‎

Lines changed: 129 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ use std::sync::Arc;
2222

2323
use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric;
2424
use crate::fuzz_cases::once_exec::OnceExec;
25-
use arrow::array::UInt64Array;
25+
use arrow::array::{AsArray, UInt64Array};
26+
use arrow::datatypes::UInt64Type;
2627
use arrow::row::{RowConverter, SortField};
2728
use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch};
2829
use arrow_schema::{DataType, Field, Schema};
@@ -34,7 +35,7 @@ use datafusion::physical_plan::sorts::sort::SortExec;
3435
use datafusion::prelude::SessionConfig;
3536
use datafusion_common::units::{KB, MB};
3637
use datafusion_execution::memory_pool::{
37-
FairSpillPool, MemoryConsumer, MemoryReservation,
38+
FairSpillPool, GreedyMemoryPool, MemoryConsumer, MemoryReservation,
3839
};
3940
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
4041
use datafusion_functions_aggregate::array_agg::array_agg_udaf;
@@ -48,7 +49,8 @@ use datafusion_physical_plan::aggregates::{
4849
use datafusion_physical_plan::metrics::MetricValue;
4950
use datafusion_physical_plan::spill::get_record_batch_memory_size;
5051
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
51-
use futures::StreamExt;
52+
use futures::{StreamExt, TryStreamExt};
53+
use rand::{Rng, SeedableRng, rngs::StdRng};
5254

5355
use arrow::array::Int32Array;
5456
use datafusion::datasource::memory::MemorySourceConfig;
@@ -832,9 +834,131 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_large_reco
832834
Ok(())
833835
}
834836

837+
#[rstest::rstest]
838+
#[case::greedy_seed_7(7, false)]
839+
#[case::greedy_seed_91(91, false)]
840+
#[case::fair_seed_7(7, true)]
841+
#[case::fair_seed_91(91, true)]
842+
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
843+
async fn aggregate_spill_with_seeded_batches_and_concurrent_peer(
844+
#[case] seed: u64,
845+
#[case] fair: bool,
846+
) -> Result<()> {
847+
const POOL_SIZE: usize = 2 * MB as usize;
848+
const INPUT_BATCHES: usize = 96;
849+
let inner: Arc<dyn MemoryPool> = if fair {
850+
Arc::new(FairSpillPool::new(POOL_SIZE))
851+
} else {
852+
Arc::new(GreedyMemoryPool::new(POOL_SIZE))
853+
};
854+
let pool = Arc::new(PeakRecordingPool::new(inner));
855+
let runtime = RuntimeEnvBuilder::new()
856+
.with_memory_pool(Arc::clone(&pool) as Arc<dyn MemoryPool>)
857+
.build_arc()?;
858+
let mut rng = StdRng::seed_from_u64(seed);
859+
let mut runs = Vec::new();
860+
// Construct both pipelines before polling either, so their reservations
861+
// overlap. Each peer reads seeded narrow and oversized string batches,
862+
// spills, and replays against the same pool with default merge fan-in.
863+
for _ in 0..2 {
864+
let batch_size = [64, 128][rng.random_range(0..2)];
865+
let sizes = Arc::new(
866+
(0..INPUT_BATCHES)
867+
.map(|index| {
868+
if index % 19 == 0 {
869+
rng.random_range(POOL_SIZE / 12..POOL_SIZE / 6)
870+
} else {
871+
rng.random_range(POOL_SIZE / 128..POOL_SIZE / 32)
872+
}
873+
})
874+
.collect::<Vec<_>>(),
875+
);
876+
let task_ctx = Arc::new(
877+
TaskContext::default()
878+
.with_session_config(SessionConfig::new().with_batch_size(batch_size))
879+
.with_runtime(Arc::clone(&runtime)),
880+
);
881+
let generated_sizes = Arc::clone(&sizes);
882+
let (args, plan, stream) =
883+
build_high_cardinality_aggregate(RunTestWithLimitedMemoryArgs {
884+
pool_size: POOL_SIZE,
885+
task_ctx,
886+
number_of_record_batches: INPUT_BATCHES,
887+
get_size_of_record_batch_to_generate: Box::pin(move |index| {
888+
generated_sizes[index]
889+
}),
890+
memory_behavior: MemoryBehavior::AsIs,
891+
assert_all_output_batches_roughly_match_batch_size_conf: false,
892+
})?;
893+
let schema = stream.schema();
894+
let mut seen = vec![false; INPUT_BATCHES * batch_size];
895+
let checked = stream.map_ok(move |batch| {
896+
let keys = batch.column(0).as_primitive::<UInt64Type>();
897+
let lists = batch.column(1).as_list::<i32>();
898+
for row in 0..batch.num_rows() {
899+
let key = keys.value(row) as usize;
900+
assert!(key < INPUT_BATCHES * batch_size, "seed={seed}, key={key}");
901+
assert!(
902+
!std::mem::replace(&mut seen[key], true),
903+
"duplicate key {key}, seed={seed}"
904+
);
905+
let expected_length = sizes[key / batch_size]
906+
.saturating_sub(size_of::<u64>() * batch_size)
907+
/ batch_size;
908+
let values = lists.value(row);
909+
let strings = values.as_string::<i32>();
910+
assert_eq!(values.len(), 1, "seed={seed}, key={key}");
911+
assert_eq!(
912+
strings.value(0).len(),
913+
expected_length,
914+
"seed={seed}, key={key}"
915+
);
916+
assert!(strings.value(0).bytes().all(|byte| byte == b'a'));
917+
}
918+
batch
919+
});
920+
let stream = Box::pin(RecordBatchStreamAdapter::new(schema, checked));
921+
runs.push(run_test(args, plan, stream, false));
922+
}
923+
tokio::time::timeout(
924+
std::time::Duration::from_secs(30),
925+
futures::future::try_join_all(runs),
926+
)
927+
.await
928+
.expect("concurrent seeded aggregates must complete")?;
929+
assert!(
930+
pool.peak_reserved() <= POOL_SIZE,
931+
"seed={seed}, fair={fair}"
932+
);
933+
assert_eq!(pool.reserved(), 0);
934+
assert_eq!(runtime.disk_manager.used_disk_space(), 0);
935+
assert_eq!(
936+
runtime.disk_manager.spilling_progress().active_files_count,
937+
0
938+
);
939+
Ok(())
940+
}
941+
835942
async fn run_test_aggregate_with_high_cardinality(
836-
mut args: RunTestWithLimitedMemoryArgs,
943+
args: RunTestWithLimitedMemoryArgs,
837944
) -> Result<MetricsSet> {
945+
let (args, aggregate_final, result) = build_high_cardinality_aggregate(args)?;
946+
// A non-spilling competitor that permanently consumes every free byte can
947+
// prevent later replay growth. Require a bounded error, not overcommit.
948+
let expect_memory_exhaustion = matches!(
949+
&args.memory_behavior,
950+
MemoryBehavior::TakeAllMemoryAtTheBeginning
951+
);
952+
run_test(args, aggregate_final, result, expect_memory_exhaustion).await
953+
}
954+
955+
fn build_high_cardinality_aggregate(
956+
mut args: RunTestWithLimitedMemoryArgs,
957+
) -> Result<(
958+
RunTestWithLimitedMemoryArgs,
959+
Arc<AggregateExec>,
960+
SendableRecordBatchStream,
961+
)> {
838962
let get_size_of_record_batch_to_generate = std::mem::replace(
839963
&mut args.get_size_of_record_batch_to_generate,
840964
Box::pin(move |_| unreachable!("should not be called after take")),
@@ -915,14 +1039,7 @@ async fn run_test_aggregate_with_high_cardinality(
9151039
)?);
9161040

9171041
let result = aggregate_final.execute(0, Arc::clone(&args.task_ctx))?;
918-
919-
// A non-spilling competitor that permanently consumes every free byte can
920-
// prevent later replay growth. Require a bounded error, not overcommit.
921-
let expect_memory_exhaustion = matches!(
922-
&args.memory_behavior,
923-
MemoryBehavior::TakeAllMemoryAtTheBeginning
924-
);
925-
run_test(args, aggregate_final, result, expect_memory_exhaustion).await
1042+
Ok((args, aggregate_final, result))
9261043
}
9271044

9281045
async fn run_test(

0 commit comments

Comments
 (0)