@@ -103,6 +103,9 @@ impl FusedStreams {
103103#[ derive( Debug ) ]
104104struct ReusableRows {
105105 inner : Vec < Option < Arc < Rows > > > ,
106+ /// Each stream's charge stays reserved while its buffer is taken for reuse.
107+ /// The sum of these sizes equals `reservation.size()`.
108+ reserved_sizes : Vec < usize > ,
106109 reservation : MemoryReservation ,
107110}
108111
@@ -130,22 +133,25 @@ impl ReusableRows {
130133 /// reservation is free for the same reason. A pool that cannot cover the buffer
131134 /// fails the query here, as it did before the buffer was cached at all.
132135 fn save ( & mut self , stream_idx : usize , rows : & Arc < Rows > ) -> Result < ( ) > {
133- self . inner [ stream_idx ] = Some ( Arc :: clone ( rows) ) ;
134- let retained = self . retained_size ( ) ;
136+ let size = rows. size ( ) ;
137+ let retained = self . reservation . size ( ) - self . reserved_sizes [ stream_idx ] + size ;
135138 debug_assert ! ( retained >= self . reservation. size( ) ) ;
136139 if let Err ( e) = self . reservation . try_resize ( retained) {
137- self . inner [ stream_idx] = None ;
138- self . reservation . try_resize ( self . retained_size ( ) ) ?;
140+ self . release ( stream_idx) ;
139141 return Err ( e) ;
140142 }
143+ self . inner [ stream_idx] = Some ( Arc :: clone ( rows) ) ;
144+ self . reserved_sizes [ stream_idx] = size;
141145 Ok ( ( ) )
142146 }
143147
144148 // drop whatever a finished stream was holding
145149 fn release ( & mut self , stream_idx : usize ) {
146150 debug_assert ! ( self . reservation. size( ) >= self . retained_size( ) ) ;
147- if let Some ( rows) = self . inner [ stream_idx] . take ( ) {
148- self . reservation . shrink ( rows. size ( ) ) ;
151+ self . inner [ stream_idx] = None ;
152+ let reserved = mem:: take ( & mut self . reserved_sizes [ stream_idx] ) ;
153+ if reserved != 0 {
154+ self . reservation . shrink ( reserved) ;
149155 }
150156 }
151157
@@ -200,6 +206,7 @@ impl RowCursorStream {
200206 streams : FusedStreams ( streams) ,
201207 rows : ReusableRows {
202208 inner : vec ! [ None ; stream_count] ,
209+ reserved_sizes : vec ! [ 0 ; stream_count] ,
203210 reservation : rows_reservation,
204211 } ,
205212 } )
@@ -442,7 +449,7 @@ mod tests {
442449 use datafusion_common:: DataFusionError ;
443450 use datafusion_execution:: RecordBatchStream ;
444451 use datafusion_execution:: memory_pool:: {
445- GreedyMemoryPool , MemoryConsumer , MemoryPool ,
452+ GreedyMemoryPool , MemoryConsumer , MemoryLimit , MemoryPool ,
446453 } ;
447454 use datafusion_physical_expr:: expressions:: col;
448455 use futures:: Stream ;
@@ -907,4 +914,36 @@ mod tests {
907914
908915 Ok ( ( ) )
909916 }
917+
918+ // The buffer has already been taken out of the cache when its reservation fails
919+ // to grow. Cleanup must still release its previous charge.
920+ #[ test]
921+ fn failed_buffer_growth_releases_only_its_stream_reservation ( ) -> Result < ( ) > {
922+ let small = ( 64 , 8 ) ;
923+ let large = ( 2048 , 128 ) ;
924+ let mut harness =
925+ RowCursorHarness :: with_shapes ( vec ! [ vec![ small, large] , vec![ small] ] ) ?;
926+
927+ // Cache stream 1's buffer and record the reservation that must survive.
928+ assert ! ( harness. poll_and_drop_cursor( 1 ) ?) ;
929+ let expected_remaining = harness. reserved ( ) ;
930+ assert ! ( harness. poll_and_drop_cursor( 0 ) ?) ;
931+
932+ // Leave no space for stream 0's larger batch.
933+ let MemoryLimit :: Finite ( limit) = harness. pool . memory_limit ( ) else {
934+ unreachable ! ( "the harness uses a bounded memory pool" ) ;
935+ } ;
936+ let competitor = MemoryConsumer :: new ( "competitor" ) . register ( & harness. pool ) ;
937+ competitor. try_grow ( limit - harness. reserved ( ) ) ?;
938+
939+ assert ! ( matches!(
940+ harness. poll_cursor( 0 ) ,
941+ Err ( DataFusionError :: ResourcesExhausted ( _) )
942+ ) ) ;
943+ drop ( competitor) ;
944+
945+ // Only stream 1's buffer and the shared converter remain reserved.
946+ assert_eq ! ( harness. reserved( ) , expected_remaining) ;
947+ Ok ( ( ) )
948+ }
910949}
0 commit comments