tustvold commented on code in PR #4522:
URL: https://github.com/apache/arrow-datafusion/pull/4522#discussion_r1051036354
##########
datafusion/core/src/execution/memory_manager/mod.rs:
##########
@@ -17,419 +17,187 @@
//! Manages all available memory during query execution
-use crate::error::{DataFusionError, Result};
-use async_trait::async_trait;
-use hashbrown::HashSet;
-use log::{debug, warn};
-use parking_lot::{Condvar, Mutex};
-use std::fmt;
-use std::fmt::{Debug, Display, Formatter};
-use std::sync::atomic::{AtomicUsize, Ordering};
+use crate::error::Result;
use std::sync::Arc;
-use std::time::{Duration, Instant};
+mod pool;
pub mod proxy;
-static CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
-
-#[derive(Debug, Clone)]
-/// Configuration information for memory management
-pub enum MemoryManagerConfig {
- /// Use the existing [MemoryManager]
- Existing(Arc<MemoryManager>),
-
- /// Create a new [MemoryManager] that will use up to some
- /// fraction of total system memory.
- New {
- /// Max execution memory allowed for DataFusion. Defaults to
- /// `usize::MAX`, which will not attempt to limit the memory
- /// used during plan execution.
- max_memory: usize,
-
- /// The fraction of `max_memory` that the memory manager will
- /// use for execution.
- ///
- /// The purpose of this config is to set aside memory for
- /// untracked data structures, and imprecise size estimation
- /// during memory acquisition. Defaults to 0.7
- memory_fraction: f64,
- },
-}
+pub use pool::*;
-impl Default for MemoryManagerConfig {
- fn default() -> Self {
- Self::New {
- max_memory: usize::MAX,
- memory_fraction: 0.7,
- }
- }
-}
+/// The pool of memory from which [`MemoryManager`] and [`TrackedAllocation`]
allocate
+pub trait MemoryPool: Send + Sync + std::fmt::Debug {
+ /// Records the creation of a new [`TrackedAllocation`] with
[`AllocationOptions`]
+ fn allocate(&self, _options: &AllocationOptions) {}
-impl MemoryManagerConfig {
- /// Create a new memory [MemoryManager] with no limit on the
- /// memory used
- pub fn new() -> Self {
- Default::default()
- }
+ /// Records the destruction of a [`TrackedAllocation`] with
[`AllocationOptions`]
+ fn free(&self, _options: &AllocationOptions) {}
- /// Create a configuration based on an existing [MemoryManager]
- pub fn new_existing(existing: Arc<MemoryManager>) -> Self {
- Self::Existing(existing)
- }
+ /// Infallibly grow the provided `allocation` by `additional` bytes
+ ///
+ /// This must always succeed
+ fn grow(&self, allocation: &TrackedAllocation, additional: usize);
- /// Create a new [MemoryManager] with a `max_memory` and `fraction`
- pub fn try_new_limit(max_memory: usize, memory_fraction: f64) ->
Result<Self> {
- if max_memory == 0 {
- return Err(DataFusionError::Plan(format!(
- "invalid max_memory. Expected greater than 0, got {}",
- max_memory
- )));
- }
- if !(memory_fraction > 0f64 && memory_fraction <= 1f64) {
- return Err(DataFusionError::Plan(format!(
- "invalid fraction. Expected greater than 0 and less than 1.0,
got {}",
- memory_fraction
- )));
- }
+ /// Infallibly shrink the provided `allocation` by `shrink` bytes
+ fn shrink(&self, allocation: &TrackedAllocation, shrink: usize);
- Ok(Self::New {
- max_memory,
- memory_fraction,
- })
- }
+ /// Attempt to grow the provided `allocation` by `additional` bytes
+ ///
+ /// On error the `allocation` will not be increased in size
+ fn try_grow(&self, allocation: &TrackedAllocation, additional: usize) ->
Result<()>;
- /// return the maximum size of the memory, in bytes, this config will allow
- fn pool_size(&self) -> usize {
- match self {
- MemoryManagerConfig::Existing(existing) => existing.pool_size,
- MemoryManagerConfig::New {
- max_memory,
- memory_fraction,
- } => (*max_memory as f64 * *memory_fraction) as usize,
- }
- }
+ /// Return the total amount of memory allocated
+ fn allocated(&self) -> usize;
}
-fn next_id() -> usize {
- CONSUMER_ID.fetch_add(1, Ordering::SeqCst)
+/// A cooperative MemoryManager which tracks memory in a cooperative fashion.
+/// `ExecutionPlan` nodes such as `SortExec`, which require large amounts of
memory
+/// register their memory requests with the MemoryManager which then tracks
the total
+/// memory that has been allocated across all such nodes.
+///
+/// The associated [`MemoryPool`] determines how to respond to memory
allocation
+/// requests, and any associated fairness control
+#[derive(Debug)]
+pub struct MemoryManager {
Review Comment:
I wonder if we even still need the `MemoryManager` :thinking:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
