diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 8cba650b79770..0d22313c07031 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -84,6 +84,7 @@ pub mod filter_pushdown; pub mod joins; pub mod limit; pub mod memory; +pub mod memory_guard; pub mod metrics; pub mod operator_statistics; pub mod placeholder_row; diff --git a/datafusion/physical-plan/src/memory_guard.rs b/datafusion/physical-plan/src/memory_guard.rs new file mode 100644 index 0000000000000..95adef472a897 --- /dev/null +++ b/datafusion/physical-plan/src/memory_guard.rs @@ -0,0 +1,26 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion_common::DataFusionError; + +#[derive(Debug)] +pub struct OperatorMemoryGuard { + consumer: MemoryConsumer, + used: AtomicUsize, +} + +impl OperatorMemoryGuard { + pub fn new(name: &str, pool: &Arc) -> Self { + let consumer = MemoryConsumer::new(name).register(pool); + Self { consumer, used: AtomicUsize::new(0) } + } + + pub fn try_reserve(&self, bytes: usize) -> Result { + if bytes == 0 { return self.consumer.try_reserve(0).map_err(|e| DataFusionError::ResourcesExhausted(e.to_string())); } + let current = self.used.load(Ordering::Relaxed); + let next = current.saturating_add(bytes); + let reservation = self.consumer.try_reserve(bytes).map_err(|e| DataFusionError::ResourcesExhausted(format!("operator memory budget exceeded: {e}")))?; + self.used.store(next, Ordering::Relaxed); + Ok(reservation) + } +} diff --git a/datafusion/physical-plan/tests/memory_guard.rs b/datafusion/physical-plan/tests/memory_guard.rs new file mode 100644 index 0000000000000..117e7bb1cb33d --- /dev/null +++ b/datafusion/physical-plan/tests/memory_guard.rs @@ -0,0 +1,4 @@ +#[test] +fn guard_rejects_when_pool_exhausted() { + // placeholder unit test skeleton; concrete pool plumbing uses GreedyMemoryPool +}