Skip to content

Commit fdaa326

Browse files
haohuaijinblaginin
andauthored
feat: impl Any for MemoryPool (#21803)
## Which issue does this PR close? - Closes #21802 ## Rationale for this change - see #21802 ## What changes are included in this PR? implment Any for MemoryPool ## Are these changes tested? yes, add one test case ## Are there any user-facing changes? --------- Co-authored-by: Dmitrii Blaginin <dmitrii@blaginin.me>
1 parent fe1dd57 commit fdaa326

2 files changed

Lines changed: 80 additions & 1 deletion

File tree

  • datafusion/execution/src/memory_pool
  • docs/source/library-user-guide/upgrading

datafusion/execution/src/memory_pool/mod.rs

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
//! help with allocation accounting.
2020
2121
use datafusion_common::{Result, internal_datafusion_err};
22+
use std::any::Any;
2223
use std::fmt::Display;
2324
use std::hash::{Hash, Hasher};
2425
use std::{cmp::Ordering, sync::Arc, sync::atomic};
@@ -182,7 +183,7 @@ pub use pool::*;
182183
///
183184
/// * [`TrackConsumersPool`]: Wraps another [`MemoryPool`] and tracks consumers,
184185
/// providing better error messages on the largest memory users.
185-
pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display {
186+
pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display {
186187
/// Return pool name
187188
fn name(&self) -> &str;
188189

@@ -224,6 +225,18 @@ pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display {
224225
}
225226
}
226227

228+
impl dyn MemoryPool {
229+
/// Returns `true` if this pool is of type `T`.
230+
pub fn is<T: MemoryPool>(&self) -> bool {
231+
(self as &dyn Any).is::<T>()
232+
}
233+
234+
/// Attempts to downcast this pool to a concrete type `T`.
235+
pub fn downcast_ref<T: MemoryPool>(&self) -> Option<&T> {
236+
(self as &dyn Any).downcast_ref()
237+
}
238+
}
239+
227240
/// Memory limit of `MemoryPool`
228241
pub enum MemoryLimit {
229242
Infinite,
@@ -603,6 +616,18 @@ mod tests {
603616
assert_eq!(pool.reserved(), 28);
604617
}
605618

619+
#[test]
620+
fn test_downcast() {
621+
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(50));
622+
623+
assert!(pool.is::<GreedyMemoryPool>());
624+
assert!(!pool.is::<UnboundedMemoryPool>());
625+
626+
let greedy: &GreedyMemoryPool = pool.downcast_ref().unwrap();
627+
assert_eq!(greedy.reserved(), 0);
628+
assert!(pool.downcast_ref::<UnboundedMemoryPool>().is_none());
629+
}
630+
606631
#[test]
607632
fn test_try_shrink() {
608633
let pool = Arc::new(GreedyMemoryPool::new(100)) as _;

docs/source/library-user-guide/upgrading/54.0.0.md

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,3 +595,57 @@ impl Default for MyTreeNode {
595595
}
596596
}
597597
```
598+
599+
### `MemoryPool` now requires `'static` (adds `Any` as a supertrait)
600+
601+
To enable downcasting of `dyn MemoryPool` to concrete pool types (via
602+
`is::<T>()` / `downcast_ref::<T>()`), the `MemoryPool` trait now has `Any`
603+
as a supertrait:
604+
605+
```rust,ignore
606+
// Before
607+
pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display { ... }
608+
609+
// After
610+
pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { ... }
611+
```
612+
613+
Because `Any` is only implemented for `'static` types, this implicitly adds a
614+
`'static` bound to every `MemoryPool` implementor.
615+
616+
**Who is affected:**
617+
618+
- Users who implement a custom `MemoryPool` whose type carries a lifetime
619+
parameter or borrows state (e.g. `struct MyPool<'a> { inner: &'a State }`).
620+
Existing implementations that are already `'static` (the common case) need
621+
no changes.
622+
623+
**Migration guide:**
624+
625+
Replace borrowed references with owned handles so the pool type becomes
626+
`'static`. The typical fix is to swap `&'a T` for `Arc<T>` (or `Rc<T>`, or an
627+
owned value):
628+
629+
```rust,ignore
630+
// Before — not 'static, no longer compiles
631+
struct MyPool<'a> {
632+
inner: &'a SomeState,
633+
}
634+
635+
impl<'a> MemoryPool for MyPool<'a> { ... }
636+
637+
// After — owned handle makes MyPool: 'static
638+
struct MyPool {
639+
inner: Arc<SomeState>,
640+
}
641+
642+
impl MemoryPool for MyPool { ... }
643+
```
644+
645+
If the borrowed state truly cannot be made `'static`, you can wrap the
646+
borrowed pool in a `'static` adapter that the pool consumer owns — for
647+
example, store the underlying state in an `Arc` owned by the adapter, or
648+
move the borrow behind an interior-mutability primitive such as `Arc<Mutex<_>>`
649+
or `Arc<OnceLock<_>>`.
650+
651+
See [PR #21803](https://github.com/apache/datafusion/pull/21803) for details.

0 commit comments

Comments
 (0)