@@ -348,6 +348,11 @@ pub(super) struct MaterializingSortMergeJoinStream {
348348 /// Tracks the number of batches currently spilled
349349 pub spilled_batch_count : usize ,
350350
351+ /// Spill-backed batches loaded for the current output window. Keep the
352+ /// original files so a later streamed row can revisit the group without
353+ /// retaining every previously read payload in memory or writing it again.
354+ restored_batches : Vec < ( usize , Arc < dyn SpillFile > ) > ,
355+
351356 /// Time spent doing the join's own work (including spill write and
352357 /// read-back). The clock is stopped while awaiting the child inputs or
353358 /// the consumer taking an emitted batch — see [`Self::stop_join_time`].
@@ -584,6 +589,7 @@ impl MaterializingSortMergeJoinStream {
584589 runtime_env,
585590 spill_manager,
586591 spilled_batch_count : 0 ,
592+ restored_batches : vec ! [ ] ,
587593 join_time,
588594 join_time_start : None ,
589595 streamed_buffered_cmp : None ,
@@ -949,7 +955,7 @@ impl MaterializingSortMergeJoinStream {
949955 fn get_required_batch_indices ( & self , buffered_freeze_count : usize ) -> Vec < usize > {
950956 let mut needed = vec ! [ ] ;
951957 // Avoid scanning if no spilled batches exist
952- if self . spilled_batch_count == 0 {
958+ if self . spilled_batch_count == 0 && self . restored_batches . is_empty ( ) {
953959 return needed;
954960 }
955961 // We need all batches that matched with streamed rows
@@ -959,9 +965,20 @@ impl MaterializingSortMergeJoinStream {
959965 }
960966 }
961967
962- // Full Joins need to emit null-joined rows, so we need batches up to freeze_count
968+ // Only batches with pending null output need to be restored here.
969+ // Restoring every batch in a Full join would load the entire key group
970+ // even when all its rows matched and none needs null materialization.
963971 if self . join_type == JoinType :: Full {
964- needed. extend ( 0 ..buffered_freeze_count) ;
972+ needed. extend (
973+ self . buffered_data
974+ . batches
975+ . iter ( )
976+ . take ( buffered_freeze_count)
977+ . enumerate ( )
978+ . filter_map ( |( idx, batch) | {
979+ ( !batch. null_joined . is_empty ( ) ) . then_some ( idx)
980+ } ) ,
981+ ) ;
965982 }
966983
967984 needed. sort_unstable ( ) ;
@@ -975,6 +992,23 @@ impl MaterializingSortMergeJoinStream {
975992 & mut self ,
976993 required_indices : & [ usize ] ,
977994 ) -> Result < ( ) > {
995+ // Drop decoded payloads outside this output's working set before
996+ // reading more. Join arrays and per-row match state remain resident.
997+ // Only visit the previous working set, not the entire buffered group.
998+ let mut retained = Vec :: new ( ) ;
999+ for ( idx, file) in std:: mem:: take ( & mut self . restored_batches ) {
1000+ if required_indices. contains ( & idx) {
1001+ retained. push ( ( idx, file) ) ;
1002+ } else {
1003+ let batch = & mut self . buffered_data . batches [ idx] ;
1004+ batch. batch = BufferedBatchState :: Spilled ( file) ;
1005+ self . reservation
1006+ . shrink ( batch. reserved_amount - batch. join_arrays_mem ) ;
1007+ batch. reserved_amount = batch. join_arrays_mem ;
1008+ self . spilled_batch_count += 1 ;
1009+ }
1010+ }
1011+ self . restored_batches = retained;
9781012 for & idx in required_indices {
9791013 // Guard against indices that might be out of bounds if the queue was cleared
9801014 if idx >= self . buffered_data . batches . len ( ) {
@@ -990,6 +1024,7 @@ impl MaterializingSortMergeJoinStream {
9901024
9911025 match spill_stream. next ( ) . await . transpose ( ) ? {
9921026 Some ( batch) => {
1027+ self . restored_batches . push ( ( idx, Arc :: clone ( spill_file) ) ) ;
9931028 // Transition the batch back to InMemory
9941029 bb. batch = BufferedBatchState :: InMemory ( batch) ;
9951030 self . spilled_batch_count -= 1 ;
@@ -1200,7 +1235,12 @@ impl MaterializingSortMergeJoinStream {
12001235 break ;
12011236 }
12021237 // load the spilled head batch before dequeuing
1203- let needed = self . get_required_batch_indices ( 1 ) ;
1238+ let mut needed = self . get_required_batch_indices ( 1 ) ;
1239+ if self . join_type == JoinType :: Full && !needed. contains ( & 0 ) {
1240+ // Deferred filter failures can create null output as the head
1241+ // is removed, even if it had no earlier null-joined indices.
1242+ needed. push ( 0 ) ;
1243+ }
12041244 self . restore_spilled_batches ( & needed) . await ?;
12051245
12061246 self . freeze_dequeuing_buffered ( ) ?;
@@ -1210,6 +1250,14 @@ impl MaterializingSortMergeJoinStream {
12101250 if matches ! ( buffered_batch. batch, BufferedBatchState :: Spilled ( _) ) {
12111251 self . spilled_batch_count -= 1 ;
12121252 }
1253+ self . restored_batches . retain_mut ( |( idx, _) | {
1254+ if * idx == 0 {
1255+ false
1256+ } else {
1257+ * idx -= 1 ;
1258+ true
1259+ }
1260+ } ) ;
12131261 head_changed = true ;
12141262 }
12151263 }
0 commit comments