DataFusion's Memory Pool: RAII Reservations, Fair Shares, and Who Gets to Spill

TL;DR: DataFusion does not instrument malloc. Instead every “big” operator registers a MemoryConsumer with a MemoryPool and holds a MemoryReservation — an RAII handle backed by an AtomicUsize that returns its bytes to the pool on Drop. try_grow is the negotiation point: success means “allocate now”, failure means “spill, shrink, or error”. FairSpillPool splits the budget between consumers that declare can_spill(true) and those that don’t, so a single operator can’t starve the rest of the plan.

Note: this tree has moved the pool out of datafusion/common/src/memory.rs. The trait and the reservation types now live in datafusion/execution/src/memory_pool/mod.rs, the concrete pools in datafusion/execution/src/memory_pool/pool.rs, and the low-level size estimation helpers in datafusion/common/src/utils/memory.rs.

The problem with counting allocations

The “obvious” design for a query engine memory limiter is to install a counting global allocator and charge every byte to the running query. It is also the design that does not work, because:

A Vec<u8> reallocating from 1 KiB to 2 KiB is charged twice unless you track capacity, not length.
Arrow’s Buffer model means two Arrays in the same RecordBatch can be zero-copy slices of the same allocation. Charging each Array double counts.
Arrow StringView/BinaryView arrays keep views into shared data buffers; slicing them retains far more memory than the slice “uses”.
The interesting question is not “who allocated this” but “which operator is holding a large working set that I could ask to shrink”.

DataFusion takes the pragmatic route, documented at length in the module header of datafusion/execution/src/memory_pool/mod.rs:49-89: streaming memory is not accounted, only the operators whose footprint grows with input rows (GroupByHashExec, SortExec, CrossJoin, …) register and constrain themselves. Everything else is “small”.

That leaves exactly three questions to answer:

How does an operator declare how much it wants?
How does the pool decide who gets it?
What happens when the answer is no?

The three types

`MemoryPool` — the policy

rust
// datafusion/execution/src/memory_pool/mod.rs:195-232
pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display {
    fn name(&self) -> &str;
    fn register(&self, _consumer: &MemoryConsumer) {}
    fn unregister(&self, _consumer: &MemoryConsumer) {}
    fn grow(&self, reservation: &MemoryReservation, additional: usize);
    fn shrink(&self, reservation: &MemoryReservation, shrink: usize);
    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()>;
    fn reserved(&self) -> usize;
    fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Unknown }
}

The trait is deliberately tiny and, crucially, the pool is told who is asking. reservation.registration.consumer is reachable from inside the pool, so a policy can make per-consumer decisions. That is the whole trick behind FairSpillPool.

grow is infallible by contract; try_grow is the one that can fail, and it must not increase the reservation on error. Failures are DataFusionError::ResourcesExhausted (pool.rs:304-318).

`MemoryConsumer` — the identity

rust
// datafusion/execution/src/memory_pool/mod.rs:270-277
pub struct MemoryConsumer {
    name: String,
    can_spill: bool,
    id: usize,
}

id is a process-unique counter (new_unique_id, mod.rs:296-300) — a consumer is not cloneable, because “the consumer” is the unit of fairness. clone_with_new_id (mod.rs:318-324) exists for the rare case where you want a second, independent registration.

The can_spill flag is the interesting one, and it is doing double duty. Its literal meaning is “this consumer can push its state to disk”. Its effective meaning, as the code comments repeatedly say, is “this consumer can handle memory back pressure”:

rust
// datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs:558-563
// We interpret 'can spill' as 'can handle memory back pressure'.
// This value needs to be set to true for the default memory pool implementations
// to ensure fair application of back pressure amongst the memory consumers.
.with_can_spill(oom_mode != OutOfMemoryMode::ReportError),

`MemoryReservation` — the RAII handle

rust
// datafusion/execution/src/memory_pool/mod.rs:382-387
pub struct MemoryReservation {
    registration: Arc<SharedRegistration>,
    size: atomic::AtomicUsize,
}

Two fields, and both earn their keep.

Arc<SharedRegistration> holds the pool and the consumer. SharedRegistration has a Drop impl that calls pool.unregister(&self.consumer) (mod.rs:370-374), so the last reservation of a consumer automatically unregisters it — the pool never sees a dangling consumer. That is why MemoryConsumer is not cloneable but MemoryReservation is cheap to derive more of: new_empty() and split() mint sibling reservations that share the same registration (mod.rs:496-522).

size: AtomicUsize is the reservation’s own byte count. grow/try_grow bump it after the pool agrees (mod.rs:472-485); free swaps it to zero and tells the pool to shrink (mod.rs:400-407).

rust
// datafusion/execution/src/memory_pool/mod.rs:525-528
impl Drop for MemoryReservation {
    fn drop(&mut self) {
        self.free();
    }
}

That Drop is the entire safety story. An operator cannot leak a reservation by returning early from an error path, by panicking, or by ?-ing its way out of a nested call. The bytes go back to the pool when the handle dies, no matter how.

The counters use Ordering::Relaxed throughout. That is safe because the pool’s own counter is the single source of truth for admission decisions, and the reservation’s size is only read for reporting (FairSpillPool does read reservation.size() inside its try_grow, under the pool mutex, but the reservation is by convention owned by one task at a time).

The three pools

`UnboundedMemoryPool`

pool.rs:33-61. A single AtomicUsize that goes up and down and never says no. It is the default, and it exists so that the accounting is still correct even when the limit is not enforced — useful for observability, and it is the fallback when a MergeMemoryPool wants to bypass the real pool (sorts/streaming_merge.rs:199-208).

`GreedyMemoryPool`

pool.rs:77-129. Classic first-come-first-served:

rust
// datafusion/execution/src/memory_pool/pool.rs:106-118
fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    self.used
        .try_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
            let new_used = used + additional;
            (new_used <= self.pool_size).then_some(new_used)
        })
        .map_err(|used| {
            insufficient_capacity_err(reservation, additional,
                self.pool_size.saturating_sub(used), self)
        })?;
    Ok(())
}

The try_update is a compare-and-swap loop, so two operators racing for the last megabyte cannot both win. Greedy is right when there is exactly one spillable operator: it hands out everything to whoever asks first, and the spiller will spill only when it truly must.

`FairSpillPool`

pool.rs:145-284. This is the interesting one, because it does not hand out memory first-come-first-served.

rust
// datafusion/execution/src/memory_pool/pool.rs:236-277
fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    let mut state = self.state.lock();
    match reservation.registration.consumer.can_spill {
        true => {
            // The total amount of memory available to spilling consumers
            let spill_available = self.pool_size.saturating_sub(state.unspillable);
            // No spiller may use more than their fraction of the memory available
            let available = spill_available
                .checked_div(state.num_spill)
                .unwrap_or(spill_available);

            if reservation.size() + additional > available {
                return Err(insufficient_capacity_err(reservation, additional, available, self));
            }
            state.spillable += additional;
        }
        false => {
            let available = self
                .pool_size
                .saturating_sub(state.unspillable + state.spillable);
            if available < additional {
                return Err(insufficient_capacity_err(reservation, additional, available, self));
            }
            state.unspillable += additional;
        }
    }
    Ok(())
}

The policy is: memory that nobody can spill is allocated greedily; the remainder is divided equally among the spillable consumers. num_spill is maintained by register/unregister (pool.rs:207-217) — the count of live spillable consumers, not a count of reservations.

The rationale is a classic livelock. Suppose a plan has two hash aggregations and a SortExec that cannot spill. With a greedy pool, the sort grabs everything, the aggregations get nothing, and both spill — writing gigabytes to disk even though the sort would have been fine with a fraction. With the fair pool, each aggregator is capped at (pool_size - unspillable) / num_spill, so both can make progress on the memory they do hold, and the sort keeps what it needs.

The doc comment is refreshingly honest about the cost:

rust
// datafusion/execution/src/memory_pool/pool.rs:145-166
/// A [`MemoryPool`] that prevents spillable reservations from using more than
/// an even fraction of the available memory sans any unspillable reservations
/// (i.e. `(pool_size - unspillable_memory) / num_spillable_reservations`)
///
/// This pool works best when you know beforehand the query has
/// multiple spillable operators that will likely all need to
/// spill. Sometimes it will cause spills even when there was
/// sufficient memory (reserved for other operators) to avoid doing
/// so.

“Sometimes it will cause spills even when there was sufficient memory” — that is the price of determinism, and it is paid on purpose.

`MergeMemoryPool`: one consumer, many reservations

A subtler problem appears when a single logical operator needs several reservations that share a budget. The merge phase and the aggregate replay phase of a spilled aggregation are two components of one operator; charging them separately would let FairSpillPool divide the operator’s own quota between its own phases.

rust
// datafusion/execution/src/memory_pool/merge_memory_pool.rs:38-52
pub struct MergeMemoryPool {
    parent: Arc<dyn MemoryPool>,
    state: Mutex<MergeMemoryState>,
}

struct MergeMemoryState {
    /// The only reservation charged to the execution pool.
    reservation: MemoryReservation,
    used: usize,
    /// Workspace already acquired for the next merge. It is reusable by any
    /// child reservation, but unavailable to other execution-pool consumers.
    retained: usize,
}

The pool charges one reservation to the parent. Children call try_grow on child reservations and the pool tops up the parent only when usage exceeds what is already reserved. retain(n) keeps workspace alive across child drops; borrow(n) hands out a WorkspaceLoan — a temporary claim on already-reserved bytes, so a merge cursor can hold a buffer without a second round trip to the parent pool (merge_memory_pool.rs:54-135).

This is the pattern the aggregate spill path uses to keep the merge buffers and the replay table under a single consumer registration (aggregates/spill.rs:273-283).

What happens when `try_grow` fails

The failure mode is entirely operator-specific, which is by design: the pool has no opinion about what “recovering” means.

Hash aggregate (aggregates/hash_stream.rs:707-750): on ResourcesExhausted, it takes the entire in-memory state out of the hash table, sorts it, writes it to a spill file, shrinks the reservation, and retries the same input loop. Only if the DiskManager has temp files disabled does it turn into a hard error: "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager".
Cross join / operators without spilling: the error propagates. The module docs spell out the difference in the three scenarios (memory_pool/mod.rs:90-113).
Repartition (repartition/mod.rs:236-247): falls back to writing the batch to a spill file instead of erroring.

Note the loop shape: it is not “reserve, fail, error”. It is “reserve, fail, release memory, retry”. The MemoryReservation API makes this ergonomic because try_resize (mod.rs:456-468) shrinks and grows in one call, and shrink/try_shrink return the freed bytes to the pool immediately.

Getting the accounting right

Two operators sharing a Buffer must not be charged twice. DataFusion solves this with a buffer-identity set:

rust
// datafusion/common/src/utils/memory.rs:152-155
pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize {
    RecordBatchMemoryCounter::new().count_batch(batch)
}

RecordBatchMemoryCounter (memory.rs:167-200) tracks start addresses of already-counted buffers, so a batch that is a slice of another batch costs only the delta. This matters directly for spilling, where the same underlying buffer can be referenced by a batch and by the state that was derived from it.

Drift: are we lying to ourselves?

MemoryDriftTracker (memory_pool/drift.rs:47-71) compares the process’s actual allocated bytes against the total reserved by all pools and logs when the difference grows. PeakRecordingPool wraps another pool and records peak usage plus the consumer that caused it.

This exists because the pool’s number is a model, not a measurement. A mis-sized estimate (estimate_memory_size, memory.rs:99-121, deliberately over-allocating buckets by 8/7 to keep hashbrown ~1/8 empty) means the pool can be off by a few percent. Drift tracking turns “we think we used 2 GiB” into “the allocator says 2.4 GiB, here is who is responsible”.

Trade-offs and failure modes

Accounting is opt-in and coarse. Memory used by streaming RecordBatches between operators is not tracked at all. The pool’s limit is a limit on state, not on RSS. Configure the pool at ~80% of your container limit, or accept the OOM-killer.
Fair division is static. FairSpillPool divides by the number of spillable consumers, not by their weight. Ten tiny spillers will halve the budget of one large one.
The can_spill flag is a promise. If a consumer claims it can spill and then cannot, the fair pool has already reserved room it will never use.
RAII hides the accounting from readers. reservation.try_grow(n) is not obviously “I am about to allocate n bytes”, and a reservation that is never shrunk keeps its bytes forever. MemoryReservation::size() and the pool’s reserved() are the tools for finding leaks; TrackConsumersPool (pool.rs:423-520) exists precisely to name the top offenders.
Relaxed atomics + Mutex state. FairSpillPool takes a parking_lot mutex on every grow/shrink. That is fine at the granularity of “one reservation change per batch”, and terrible if you call it per row.

The upshot is a system that is not trying to be a general-purpose allocator accountant. It is a quota system with RAII: operators ask, the pool grants or denies, and denial is a normal, expected control-flow event rather than an error. That is the design decision that makes spilling, sorting and joins possible at all — because in a system where memory denial is exceptional, operators must be written to expect it.

The problem with counting allocations

The three types

MemoryPool — the policy

MemoryConsumer — the identity

MemoryReservation — the RAII handle