Do not require mut in memory reservation methods #19759

gabotechs · 2026-01-12T08:48:11Z

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.

alamb

Looks good to me -- thank you @gabotechs

alamb · 2026-01-23T18:48:57Z

datafusion/execution/src/memory_pool/mod.rs

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


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)

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.

alamb · 2026-01-23T18:50:23Z

datafusion/execution/src/memory_pool/mod.rs

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


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

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.

Running some benchmarks just in case

We should be good

gabotechs · 2026-01-26T08:46:44Z

run benchmark tpch

alamb-ghbot · 2026-01-26T08:46:53Z

🤖 ./gh_compare_branch.sh gh_compare_branch.sh Running
Linux aal-dev 6.14.0-1018-gcp #19~24.04.1-Ubuntu SMP Wed Sep 24 23:23:09 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing do-not-require-mut-in-memory-reservation (4f560aa) to ca904b3 diff using: tpch
Results will be posted here when complete

alamb-ghbot · 2026-01-26T09:02:50Z

🤖: Benchmark completed

Details

Comparing HEAD and do-not-require-mut-in-memory-reservation
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query     ┃      HEAD ┃ do-not-require-mut-in-memory-reservation ┃        Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1  │ 181.88 ms │                                179.41 ms │     no change │
│ QQuery 2  │  93.11 ms │                                 90.75 ms │     no change │
│ QQuery 3  │ 126.73 ms │                                124.61 ms │     no change │
│ QQuery 4  │  78.68 ms │                                 77.48 ms │     no change │
│ QQuery 5  │ 172.91 ms │                                176.08 ms │     no change │
│ QQuery 6  │  65.44 ms │                                 67.89 ms │     no change │
│ QQuery 7  │ 208.70 ms │                                212.51 ms │     no change │
│ QQuery 8  │ 168.46 ms │                                171.42 ms │     no change │
│ QQuery 9  │ 236.93 ms │                                229.86 ms │     no change │
│ QQuery 10 │ 187.68 ms │                                189.28 ms │     no change │
│ QQuery 11 │  64.89 ms │                                 62.80 ms │     no change │
│ QQuery 12 │ 117.70 ms │                                118.81 ms │     no change │
│ QQuery 13 │ 219.34 ms │                                218.82 ms │     no change │
│ QQuery 14 │  92.66 ms │                                 95.56 ms │     no change │
│ QQuery 15 │ 122.08 ms │                                127.40 ms │     no change │
│ QQuery 16 │  60.98 ms │                                 61.68 ms │     no change │
│ QQuery 17 │ 261.15 ms │                                265.50 ms │     no change │
│ QQuery 18 │ 313.24 ms │                                315.39 ms │     no change │
│ QQuery 19 │ 138.29 ms │                                138.34 ms │     no change │
│ QQuery 20 │ 133.04 ms │                                133.18 ms │     no change │
│ QQuery 21 │ 262.93 ms │                                256.93 ms │     no change │
│ QQuery 22 │  43.46 ms │                                 41.14 ms │ +1.06x faster │
└───────────┴───────────┴──────────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary                                       ┃           ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (HEAD)                                       │ 3350.30ms │
│ Total Time (do-not-require-mut-in-memory-reservation)   │ 3354.82ms │
│ Average Time (HEAD)                                     │  152.29ms │
│ Average Time (do-not-require-mut-in-memory-reservation) │  152.49ms │
│ Queries Faster                                          │         1 │
│ Queries Slower                                          │         0 │
│ Queries with No Change                                  │        21 │
│ Queries with Failure                                    │         0 │
└─────────────────────────────────────────────────────────┴───────────┘

alamb · 2026-01-26T11:44:42Z

Thanks @gabotechs

## Which issue does this PR close?  - Closes #. ## Rationale for this change This is a PR from a batch of PRs that attempt to improve performance in hash joins: - #19759 - This PR - #19761 It adds a building block that allows eagerly collecting data on the probe side of a hash join before the build side is finished. Even if the intended use case is for hash joins, the new execution node is generic and is designed to work anywhere in the plan. ## What changes are included in this PR? > [!NOTE] > The new BufferExec node introduced in this PR is still not wired up automatically  Adds a new `BufferExec` node that can buffer up to a certain size in bytes for each partition eagerly performing work that otherwise would be delayed. Schematically, it looks like this: ``` ┌───────────────────────────┐ │ BufferExec │ │ │ │┌────── Partition 0 ──────┐│ ││ ┌────┐ ┌────┐││ ┌────┐ ──background poll────────▶│ │ │ ├┼┼───────▶ │ ││ └────┘ └────┘││ └────┘ │└─────────────────────────┘│ │┌────── Partition 1 ──────┐│ ││ ┌────┐ ┌────┐ ┌────┐││ ┌────┐ ──background poll─▶│ │ │ │ │ ├┼┼───────▶ │ ││ └────┘ └────┘ └────┘││ └────┘ │└─────────────────────────┘│ │ │ │ ... │ │ │ │┌────── Partition N ──────┐│ ││ ┌────┐││ ┌────┐ ──background poll───────────────▶│ ├┼┼───────▶ │ ││ └────┘││ └────┘ │└─────────────────────────┘│ └───────────────────────────┘ ``` ## Are these changes tested?  yes, by new unit tests ## Are there any user-facing changes?  users can import a new `BufferExec` execution plan in their codebase, but no internal usage is shipped yet in this PR.

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

This was referenced Jan 12, 2026

Hash join buffering on probe side #19761

Open

Add BufferExec execution plan #19760

Merged

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

Do not require mut in memory reservation methods

4f560aa

gabotechs force-pushed the do-not-require-mut-in-memory-reservation branch from a578d54 to 4f560aa Compare January 16, 2026 10:29

gabotechs marked this pull request as ready for review January 19, 2026 14:49

alamb approved these changes Jan 24, 2026

View reviewed changes

alamb added this pull request to the merge queue Jan 26, 2026

alamb added the api change Changes the API exposed to users of the crate label Jan 26, 2026

Merged via the queue into apache:main with commit b6f7521 Jan 26, 2026
32 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Do not require mut in memory reservation methods #19759

Do not require mut in memory reservation methods #19759

Uh oh!

gabotechs commented Jan 12, 2026 •

edited

Loading

Uh oh!

alamb left a comment

Uh oh!

alamb Jan 23, 2026

Uh oh!

gabotechs Jan 24, 2026

Uh oh!

alamb Jan 23, 2026

Uh oh!

gabotechs Jan 24, 2026

Uh oh!

gabotechs Jan 26, 2026

Uh oh!

gabotechs Jan 26, 2026

Uh oh!

gabotechs commented Jan 26, 2026

Uh oh!

alamb-ghbot commented Jan 26, 2026

Uh oh!

alamb-ghbot commented Jan 26, 2026

Uh oh!

alamb commented Jan 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Do not require mut in memory reservation methods #19759

Do not require mut in memory reservation methods #19759

Uh oh!

Conversation

gabotechs commented Jan 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

alamb Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs Jan 24, 2026

Choose a reason for hiding this comment

Uh oh!

alamb Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs Jan 24, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs commented Jan 26, 2026

Uh oh!

alamb-ghbot commented Jan 26, 2026

Uh oh!

alamb-ghbot commented Jan 26, 2026

Uh oh!

alamb commented Jan 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

gabotechs commented Jan 12, 2026 •

edited

Loading