Under reservation in multi level merge due to extra buffering in spawn_buffered

Describe the bug

Basically in multi level merge:

datafusion/datafusion/physical-plan/src/sorts/multi_level_merge.rs

Lines 456 to 465 in b7e7b51

    
           // Free the memory we reserved for this merge as we either try again or fail 
        
           reservation.free(); 
        
           if buffer_len > 1 { 
        
               // Try again with smaller buffer size, it will be slower but at least we can merge 
        
               return self.get_sorted_spill_files_to_merge( 
        
                   buffer_len - 1, 
        
                   minimum_number_of_required_streams, 
        
                   reservation, 
        
               ); 
        
           }

the reservation is multiplied by buffer_len.

This buffer_size is later on used to read the spill files as a stream:

datafusion/datafusion/physical-plan/src/sorts/multi_level_merge.rs

Lines 329 to 339 in b7e7b51

    
           for spill in sorted_spill_files { 
        
               let stream = self 
        
                   .spill_manager 
        
                   .clone() 
        
                   .with_batch_read_buffer_capacity(buffer_size) 
        
                   .read_spill_as_stream( 
        
                       spill.file, 
        
                       Some(spill.max_record_batch_memory), 
        
                   )?; 
        
               sorted_streams.push(stream); 
        
           }

read_spill_as_stream calls spawn_buffered with self.batch_read_buffer_capacity:

datafusion/datafusion/physical-plan/src/spill/spill_manager.rs

Line 194 in b7e7b51

Ok(spawn_buffered(stream, self.batch_read_buffer_capacity))

The maximum amount of RecordBatches produced by this spawn_buffered call:

datafusion/datafusion/physical-plan/src/common.rs

Lines 191 to 208 in 739e147

    
           Ok(handle) 
        
               if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => 
        
           { 
        
               let mut builder = RecordBatchReceiverStream::builder(input.schema(), buffer); 
        
               let sender = builder.tx(); 
        
               builder.spawn(async move { 
        
                   while let Some(item) = input.next().await { 
        
                       if sender.send(item).await.is_err() { 
        
                           // Receiver dropped when query is shutdown early (e.g., limit) or error, 
        
                           // no need to return propagate the send error. 
        
                           return Ok(()); 
        
                       } 
        
                   } 
        
                   Ok(()) 
        
               });

is actually buffer + 1.
Example with buffer = 1:
RecordBatch is produced, channel capacity is 1, sender.send succeeds, 1 RecordBatch in memory
RecordBatch is produced, channel is full, sender.send blocks, but the RecordBatch has already been produced => 2 RecordBatches in memory

This is an issue especially with large Record batches

To Reproduce

No response

Expected behavior

Since buffer needs to be at least one, get a permit before producing the RecordBatch: https://docs.rs/tokio/latest/tokio/sync/mpsc/struct.Permit.html
Update all the callers to use buffer + 1 in order to preserve behavior?

Additional context

No response

take

fix: ensure a maximum of `buffer_len` RecordBatches are cached in `sp…

Navigation Menu

Repository navigation

Under reservation in multi level merge due to extra buffering in spawn_buffered #23559

Describe the bug

To Reproduce

Expected behavior

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	// Free the memory we reserved for this merge as we either try again or fail
	reservation.free();
	if buffer_len > 1 {
	// Try again with smaller buffer size, it will be slower but at least we can merge
	return self.get_sorted_spill_files_to_merge(
	buffer_len - 1,
	minimum_number_of_required_streams,
	reservation,
	);
	}

	for spill in sorted_spill_files {
	let stream = self
	.spill_manager
	.clone()
	.with_batch_read_buffer_capacity(buffer_size)
	.read_spill_as_stream(
	spill.file,
	Some(spill.max_record_batch_memory),
	)?;
	sorted_streams.push(stream);
	}

	Ok(handle)
	if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread =>
	{
	let mut builder = RecordBatchReceiverStream::builder(input.schema(), buffer);

	let sender = builder.tx();

	builder.spawn(async move {
	while let Some(item) = input.next().await {
	if sender.send(item).await.is_err() {
	// Receiver dropped when query is shutdown early (e.g., limit) or error,
	// no need to return propagate the send error.
	return Ok(());
	}
	}

	Ok(())
	});

Under reservation in multi level merge due to extra buffering in spawn_buffered #23559

Description

Describe the bug

To Reproduce

Expected behavior

Additional context

Activity

ariel-miculas commented on Jul 14, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions