Skip to content

Conversation

@gabotechs
Copy link
Contributor

@gabotechs gabotechs commented Jan 12, 2026

Which issue does this PR close?

  • Closes #.

Rationale for this change

Prerequisite for the following PRs:

Even if the api on the MemoryPool does not require &mut self for growing/shrinking the reserved size, the api in MemoryReservation does, making simple implementations irrepresentable without synchronization primitives. For example, the following would require a Mutex for concurrent access to the MemoryReservation in different threads, even though the MemoryPool doesn't:

let mut stream: SendableRecordBatchStream = SendableRecordBatchStream::new();
let mem: Arc<MemoryReservation> = Arc::new(MemoryReservation::new_empty());

let mut builder = ReceiverStreamBuilder::new(10);
let tx = builder.tx();
{
    let mem = mem.clone();
    builder.spawn(async move {
        while let Some(msg) = stream.next().await {
            mem.try_grow(msg.unwrap().get_array_memory_size()); // ❌ `mem` is not mutable
            tx.send(msg).unwrap();
        }
    });
}
builder
    .build()
    .inspect_ok(|msg| mem.shrink(msg.get_array_memory_size()));  // ❌ `mem` is not mutable

What changes are included in this PR?

Make the methods in MemoryReservation require &self instead of &mut self for allowing concurrent shrink/grows from different tasks for the same reservation.

Are these changes tested?

yes, by current tests

Are there any user-facing changes?

Users can now safely call methods of MemoryReservation from different tasks without synchronization primitives.

This is a backwards compatible API change, as it will work out of the box for current users, however, depending on their clippy configuration, they might see some new warnings about "unused muts" in their codebase.

@github-actions github-actions bot added core Core DataFusion crate execution Related to the execution crate datasource Changes to the datasource crate physical-plan Changes to the physical-plan crate labels Jan 12, 2026
@gabotechs gabotechs force-pushed the do-not-require-mut-in-memory-reservation branch 2 times, most recently from 2a538b1 to a578d54 Compare January 13, 2026 16:51
@gabotechs gabotechs force-pushed the do-not-require-mut-in-memory-reservation branch from a578d54 to 4f560aa Compare January 16, 2026 10:29
@gabotechs gabotechs marked this pull request as ready for review January 19, 2026 14:49
Copy link
Contributor

@alamb alamb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me -- thank you @gabotechs

/// pool, returning the number of bytes freed.
pub fn free(&mut self) -> usize {
let size = self.size;
pub fn free(&self) -> usize {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this technically this is a breaking API change? I thought about it and from what I can tell the answer is no as to all this API in DataFusion 52 the caller needs a mut and in 53 would not (but could still call it with mut even though that is not needed)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's right, this is not a breaking API change.

One thing that could happen is that people have some clippy lint that goes off in case of "unused muts". In that case people will start seeing new clippy warnings with DataFusion 53 in their own code.

pub fn free(&mut self) -> usize {
let size = self.size;
pub fn free(&self) -> usize {
let size = self.size.load(atomic::Ordering::Relaxed);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would expect that the reservations should be consistent (Ordering::seqcst), otherwise I would worry that we run the risk of not seeing other changes.

However, Relaxed seems to be used in the MemoryPools themselves, so this is consistent

https://github.com/apache/datafusion/blob/ead8209803770773980fafaf0fc622bb606be0ee/datafusion/execution/src/memory_pool/pool.rs#L83-L82

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, the actual accounting that matters is done in the impl MemoryPools, so using something more consistent here is not really going to yield any improvement.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate datasource Changes to the datasource crate execution Related to the execution crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants