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