diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index 829e313d2381e..2b36ee7f40add 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -19,6 +19,7 @@ //! help with allocation accounting. use datafusion_common::{Result, internal_datafusion_err}; +use std::any::Any; use std::fmt::Display; use std::hash::{Hash, Hasher}; use std::{cmp::Ordering, sync::Arc, sync::atomic}; @@ -182,7 +183,7 @@ pub use pool::*; /// /// * [`TrackConsumersPool`]: Wraps another [`MemoryPool`] and tracks consumers, /// providing better error messages on the largest memory users. -pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display { +pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { /// Return pool name fn name(&self) -> &str; @@ -224,6 +225,18 @@ pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display { } } +impl dyn MemoryPool { + /// Returns `true` if this pool is of type `T`. + pub fn is(&self) -> bool { + (self as &dyn Any).is::() + } + + /// Attempts to downcast this pool to a concrete type `T`. + pub fn downcast_ref(&self) -> Option<&T> { + (self as &dyn Any).downcast_ref() + } +} + /// Memory limit of `MemoryPool` pub enum MemoryLimit { Infinite, @@ -603,6 +616,18 @@ mod tests { assert_eq!(pool.reserved(), 28); } + #[test] + fn test_downcast() { + let pool: Arc = Arc::new(GreedyMemoryPool::new(50)); + + assert!(pool.is::()); + assert!(!pool.is::()); + + let greedy: &GreedyMemoryPool = pool.downcast_ref().unwrap(); + assert_eq!(greedy.reserved(), 0); + assert!(pool.downcast_ref::().is_none()); + } + #[test] fn test_try_shrink() { let pool = Arc::new(GreedyMemoryPool::new(100)) as _; diff --git a/docs/source/library-user-guide/upgrading/54.0.0.md b/docs/source/library-user-guide/upgrading/54.0.0.md index 34d1f7c61eaf1..9ce30c97bba97 100644 --- a/docs/source/library-user-guide/upgrading/54.0.0.md +++ b/docs/source/library-user-guide/upgrading/54.0.0.md @@ -497,3 +497,57 @@ impl Default for MyTreeNode { } } ``` + +### `MemoryPool` now requires `'static` (adds `Any` as a supertrait) + +To enable downcasting of `dyn MemoryPool` to concrete pool types (via +`is::()` / `downcast_ref::()`), the `MemoryPool` trait now has `Any` +as a supertrait: + +```rust,ignore +// Before +pub trait MemoryPool: Send + Sync + std::fmt::Debug + Display { ... } + +// After +pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { ... } +``` + +Because `Any` is only implemented for `'static` types, this implicitly adds a +`'static` bound to every `MemoryPool` implementor. + +**Who is affected:** + +- Users who implement a custom `MemoryPool` whose type carries a lifetime + parameter or borrows state (e.g. `struct MyPool<'a> { inner: &'a State }`). + Existing implementations that are already `'static` (the common case) need + no changes. + +**Migration guide:** + +Replace borrowed references with owned handles so the pool type becomes +`'static`. The typical fix is to swap `&'a T` for `Arc` (or `Rc`, or an +owned value): + +```rust,ignore +// Before — not 'static, no longer compiles +struct MyPool<'a> { + inner: &'a SomeState, +} + +impl<'a> MemoryPool for MyPool<'a> { ... } + +// After — owned handle makes MyPool: 'static +struct MyPool { + inner: Arc, +} + +impl MemoryPool for MyPool { ... } +``` + +If the borrowed state truly cannot be made `'static`, you can wrap the +borrowed pool in a `'static` adapter that the pool consumer owns — for +example, store the underlying state in an `Arc` owned by the adapter, or +move the borrow behind an interior-mutability primitive such as `Arc>` +or `Arc>`. + +See [PR #21803](https://github.com/apache/datafusion/pull/21803) for details.