feat: Add unique id for every memory consumer by EmilyMatt · Pull Request #15613 · apache/datafusion

EmilyMatt · 2025-04-07T09:15:46Z

Which issue does this PR close?

Closes Unique identifier for MemoryConsumer #15126 .

Rationale for this change

This makes an important change to the MemoryConsumer API, letting any memory pool identify which consumer is requesting the memory, and therefor separate consumers that might have the same name (for example, two AggregateExec operators in the same Partition)
This also modifies the TrackConsumer pool to adhere to these changes.

What changes are included in this PR?

Nothing except the addition of a unique identifier in the consumer, and an id() function to retrieve it.

Are these changes tested?

All of the TrackedConsumer pool tests now verify this, as they have all been updated to accomodate the fact that the pool uses the IDs to track (the inner pool is unchanged in its behaviour atm)
This ensures the validity of the logic

Are there any user-facing changes?

Yes, MemoryConsumer is no longer clone-able, it is now moved into the register() function and thats it, to avoid an invalid state where a MemoryConsumer can be registered more than once.
The TrackedConsumerMemoryPool will now store Consumers based on their IDs, instead of names, so will be able to track multiple consumers with the same name, based on their ID, and the error messages are updated accordingly.

alamb

Thank you @EmilyMatt -- this API make a lot of sense to me

I think it would be useful to add tests for this feature, so that we don't break it in the future accidentally

Specifically, I think testing the "memory consumer ids are unique" behavior would be best as that is the goal of the API. SO for example, perhaps write a query that has multiple AggregateExecs and verify that the assigned memory consumer Ids are different

I would also recommend adding ids to the various messages (such as in TrackedConsumerPool) https://github.com/apache/datafusion/blob/main/datafusion/execution/src/memory_pool/pool.rs#L261

So that the id can be correlated in more places

EmilyMatt · 2025-04-07T13:09:55Z

Thank you @EmilyMatt -- this API make a lot of sense to me

I think it would be useful to add tests for this feature, so that we don't break it in the future accidentally

Specifically, I think testing the "memory consumer ids are unique" behavior would be best as that is the goal of the API. SO for example, perhaps write a query that has multiple AggregateExecs and verify that the assigned memory consumer Ids are different

I would also recommend adding ids to the various messages (such as in TrackedConsumerPool) https://github.com/apache/datafusion/blob/main/datafusion/execution/src/memory_pool/pool.rs#L261

So that the id can be correlated in more places

I have modified this PR to also reflect these changes in the TrackConsumer pool, but in order to ensure a valid state this also requires a breaking API change - MemoryConsumer no longer implements Hash and Clone
This is in order to ensure a user does not call clone() on a consumer, and then register the cloned one, while still maintaining an owned MemoryConsumer, which can be called with register() again, despite having the same id as an already existing MemoryConsumer in the pool.

I'd have preferred to deprecate the Clone implementation and only removing it in a future version, along with an explanation, but unfortunately, traits cannot be deprecated using the attribute.
This means this PR needs the api change label, but I can't manually set it^

berkaysynnada

Are we planning to deprecate MemoryConsumer and replace it all with TrackedConsumer ?

EmilyMatt · 2025-04-08T14:56:23Z

Are we planning to deprecate MemoryConsumer and replace it all with TrackedConsumer ?

I think TrackedConsumer is only relevant for the TrackConsumerPool, it doesn't exist anywhere else, most pools don't track consumers at all, it is up to users to implement a MemoryPool that takes into consideration each Consumer and how much it has requested.
This API is only meant to help users do that, as the current API does not really allow doing that with certainty(there can be surrounding hacks to achieve that functionality, but they far from perfect without some unique identifier)

alamb

Thank you @EmilyMatt -- I think this looks like a very nice improvement to me

FYI @wiedld and @2010YOUY01

I also marked this PR as an API change as MemoryConsumer no longer implements Clone.

alamb · 2025-04-09T14:51:48Z

+        }
+    }
+
+    pub fn clone_with_new_id(&self) -> Self {


Suggested change

pub fn clone_with_new_id(&self) -> Self {

/// Return a clone of the MemoryConsumer that has a new

/// unique ID assigned.

pub fn clone_with_new_id(&self) -> Self {

No problem, added some more documentation in general

alamb · 2025-04-09T14:52:55Z

 ///
 /// [proxy]: datafusion_common::utils::proxy
-#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+#[derive(Debug)]


Can you also please add a note to this doc string that says each MemoryConsumer gets an automatically assigned globally unique id and that to clone a memory consumer, see [MemoryConsumer::clone_with_new_id]

alamb · 2025-04-09T14:55:02Z

+                        tracked_consumer.name.to_owned(),
+                        tracked_consumer.can_spill,
+                    ),
+                    tracked_consumer.reserved.load(Ordering::Acquire),


As a follow on / cleanup it would be possible to refactor this code into methods, so like this (and it might make the code easier to read)

Suggested change

tracked_consumer.reserved.load(Ordering::Acquire),

tracked_consumer.reserved(),

The same thing goes for other updates later down

Sounds good to me, done^

alamb · 2025-04-09T14:55:45Z

        // using the previously set sizes for other consumers
        let mut r5 = MemoryConsumer::new("r5").register(&pool);
-        let expected = "Additional allocation failed with top memory consumers (across reservations) as: r1 consumed 50 bytes, r3 consumed 20 bytes, r2 consumed 15 bytes. Error: Failed to allocate additional 150 bytes for r5 with 0 bytes already allocated for this reservation - 5 bytes remain available for the total pool";
+        let expected = format!("Additional allocation failed with top memory consumers (across reservations) as: r1#{}(can spill: false) consumed 50 bytes, r3#{}(can spill: false) consumed 20 bytes, r2#{}(can spill: false) consumed 15 bytes. Error: Failed to allocate additional 150 bytes for r5 with 0 bytes already allocated for this reservation - 5 bytes remain available for the total pool", r1.consumer().id(), r3.consumer().id(), r2.consumer().id());


this is a nice change

EmilyMatt · 2025-04-09T19:00:57Z

@alamb Thanks for taking the time to look at this. I've addressed the relevant points and added some general documentation, I've also modified the reserved field to an AtomicUsize(instead of AtomicU64), as that feels more inline with the memory pool API, which uses usize everywhere, I hope that's not overtly obtrusive.

alamb · 2025-04-11T20:20:41Z

This is great -- thanks again for the work and contribution @EmilyMatt

* feat: Add unique id for every memory consumer * imports * feat: Make MemoryConsumer non-clonable, update tracked pool to use ids * Add hash impl and fmt * use a single cfg instead wrapping the if * Address review, add documentation

EmilyMatt added 2 commits April 7, 2025 11:14

feat: Add unique id for every memory consumer

68cbddd

imports

dcb2069

github-actions Bot added the execution Related to the execution crate label Apr 7, 2025

alamb reviewed Apr 7, 2025

View reviewed changes

feat: Make MemoryConsumer non-clonable, update tracked pool to use ids

d01b3d9

alamb mentioned this pull request Apr 7, 2025

Weekly Plan (Andrew Lamb) April 7, 2025 #15616

Closed

12 tasks

alamb added the api change Changes the API exposed to users of the crate label Apr 7, 2025

alamb reviewed Apr 7, 2025

View reviewed changes

Comment thread datafusion/execution/src/memory_pool/mod.rs

EmilyMatt added 2 commits April 7, 2025 16:23

Add hash impl and fmt

75bbbb3

use a single cfg instead wrapping the if

2f3f9e5

berkaysynnada reviewed Apr 8, 2025

View reviewed changes

alamb approved these changes Apr 9, 2025

View reviewed changes

Address review, add documentation

d28c9d8

alamb merged commit 263222e into apache:main Apr 11, 2025

	tracked_consumer.reserved.load(Ordering::Acquire),
	tracked_consumer.reserved(),

Conversation

EmilyMatt commented Apr 7, 2025 • 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!

EmilyMatt commented Apr 7, 2025

Uh oh!

Uh oh!

berkaysynnada left a comment

Choose a reason for hiding this comment

Uh oh!

EmilyMatt commented Apr 8, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EmilyMatt commented Apr 9, 2025

Uh oh!

alamb commented Apr 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

EmilyMatt commented Apr 7, 2025 •

edited

Loading