External sort failing with non-spillable operators as input (RepartitionExec)

[16pierre]

Describe the bug

I'm trying to sort some input (union of filtered Parquet files) with fixed parallelism, in order to do so I'm manually fiddling with a round-robin RepartitionExec operator before the sort. I reckon this is slightly hacky, they may be alternatives, the possibility of unions of large numbers of inputs makes partitioning trickier though.

When doing so, I get memory allocation failures on the RepartitionExec operator (I suspect this could happen with other non-spillable operators used as input of the Sort nodes though):

ResourcesExhausted(\"Additional allocation failed with top memory consumers (across reservations) as: 
ExternalSorter[0] consumed 145849856 bytes, 
ExternalSorterMerge[0] consumed 50000000 bytes,
RepartitionExec[0] consumed 2951744 bytes. 

Error: Failed to allocate additional 2951744 bytes for RepartitionExec[0] with 2951744 bytes already allocated for this reservation - 1198400 bytes remain available for the total pool

If my understanding is correct, this feels expected given the current sort implementation: we trigger spill only after saturating the memory pool and getting allocation failure on try_grow:

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Line 787 in 033fc9b

match self.reservation.try_grow(size) {

Which means that once the global memory pool is filled up with sort allocations, any operator used as input of the sort can fail if it uses non-spillable memory allocations.

To Reproduce

My current repro setup cannot be shared as a snippet because it uses a bunch of production codepaths.

Roughly what I'm doing is:

• write a Parquet file with a couple millions rows (binary columns, 16 bytes each, UUIDs)
• read the Parquet file and run repartition + sort

// Memory allocator = FairSpillPool

let physical_plan = open_my_parquet_file(...);

// Note: I get failures with input partitions = 1 and sort_parallelism = 1
let partitioning = Partitioning::RoundRobinBatch(sort_parallelism as usize);
let repartitioned_exec = Arc::new(RepartitionExec::try_new(physical_plan, partitioning))
SortExec::new(sort_expression.clone(), repartitioned_exec)

Note we're currently on Datafusion 46.0.1 (will upgrade soon to pick up recent sort improvements)

Expected behavior

No response

Additional context

Main work-arounds/hacks that come to mind:

• register a fake spillable allocation to force FairSpillPool to spill earlier than before memory is full
• implement custom MemoryPool to special-case how different operators are tracked

Ideally I'd prefer a solution where we can precisely control how much memory is allocated to the sort compared to the rest of the operators, especially given the subtle memory allocations we faced during the sort merge phase - the solution "fake spillable allocation" is a bit coarse-grain with that regards.

[2010YOUY01]

The current FairSpillPool implementation seems problematic
See its behavior:

datafusion/datafusion/execution/src/memory_pool/pool.rs

Line 202 in 5021b39

fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {

I think some specific timing can cause the non-spillable operators to OOM, however if spillable operators can spill earlier, the execution should be possible to finish. (@ding-young Just want to make sure, is such race condition possible? 🤔 )

• implement custom MemoryPool to special-case how different operators are tracked

Ideally I'd prefer a solution where we can precisely control how much memory is allocated to the sort compared to the rest of the operators, especially given the subtle memory allocations we faced during the sort merge phase - the solution "fake spillable allocation" is a bit coarse-grain with that regards.

This fine-grained memory poll approach sounds good. If we put all non-spillable consumers (like RepartitionExec) to one pool, and all spillable consumers (like SortExec) to another pool, such race condition won't be possible.

Besides I think maybe it's okay to remove memory reservations in RepartitionExec, it's memory usage should be bounded by batch_mem_size * partition_count, which is a constant.

---

BTW spilling execution is still in experimental stage, we're working on bug-fixes to make it production ready. If you can provide a runnable reproducer, we can help further investigate.

[16pierre]

Thanks a lot for the quick feedback !

I'm a bit stretch on dev bandwidth as we're still in early phases of our Datafusion production integration, but will try my best to fork our code and put up a reproducer.

[16pierre]

I think some specific timing can cause the non-spillable operators to OOM, however if spillable operators can spill earlier, the execution should be possible to finish

I could definitely foresee race-conditions being possible, but I can also imagine deterministic failures. For example in case of a single partition: we deterministically fill up memory with the batches to sort, and let's say we get to the point where we need one more batch to trigger sort spilling. At this point the non-spillable input operator needs to allocate memory to provide that last batch, and errors because not enough memory available.

[ding-young]

I think some specific timing can cause the non-spillable operators to OOM, however if spillable operators can spill earlier, the execution should be possible to finish. (@ding-young Just want to make sure, is such race condition possible? 🤔 )

Yes, scheduling / timing matters here definitely since we don't have any mechanism to reclaim memory reservation from other spillable operators when non-spillable operators fail to try_grow its reservation.

I could definitely foresee race-conditions being possible, but I can also imagine deterministic failures. For example in case of a single partition: we deterministically fill up memory with the batches to sort, and let's say we get to the point where we need one more batch to trigger sort spilling. At this point the non-spillable input operator needs to allocate memory to provide that last batch, and errors because not enough memory available.

But yes, this is also possible.

[16pierre]

Besides I think maybe it's okay to remove memory reservations in RepartitionExec, it's memory usage should be bounded by batch_mem_size * partition_count, which is a constant.

Interestingly, when experimenting locally with workarounds suggested in this issue (splitting the memory pools), I've been hitting OOM issues where RepartitionExec was consuming much more than batch_mem_size * partition_count (FYI I'm using RoundRobin)

Still investigating, I suspect that upon each RepartitionExec#pull_from_input loop iteration, we push one RecordBatch to each candidate output_channels, no matter how many batches are currently buffered ? See these codepaths:

datafusion/datafusion/physical-plan/src/repartition/mod.rs

Lines 934 to 935 in 4bc66c8

	// While there are still outputs to send to, keep pulling inputs
	let mut batches_until_yield = partitioner.num_partitions();

datafusion/datafusion/physical-plan/src/repartition/mod.rs

Line 957 in 4bc66c8

if tx.send(Some(Ok(batch))).await.is_err() {

At first glance of the tx.send codepath I'm not seeing clear hints that this enforces bounded capacity and can back-pressure - which could make sense, I guess that blocking on this call could lead to all kinds of multi-threading/dead-lock problems, but I may be wrong here, this codepath isn't trivial. Docs seem to point in this direction too:

datafusion/datafusion/physical-plan/src/repartition/distributor_channels.rs

Lines 37 to 39 in e718c1a

	//! There are `N` virtual MPSC (multi-producer, single consumer) channels with unbounded capacity. However, if all
	//! buffers/channels are non-empty, than a global gate will be closed preventing new data from being written (the
	//! sender futures will be [pending](Poll::Pending)) until at least one channel is empty (and not closed).

I interpret this as: if there is unbalance between the partition processing speeds, the RepartitionExec node could buffer a bunch of RecordBatches on some partitions, at least temporarily. And sorting could very well be a big offender here because some partitions get naturally stuck for long durations when they need to sort-spill.

If this understanding is correct maybe I should consider controlling the parallelism differently than with a RepartitionExec here.

[milenkovicm]

FWIW, we take an optimistic approach; we do not limit memory for non-spillable operators. Getting numbers right when non-spillable operators are accounted for was impossible for us. Actually, we do track the memory usage of them, but we let it grow unbounded.

The other reason is that we can’t do much with them; from our experience, they freak out a bit too early before a spill can occur.

Also, this makes the estimation of overall process memory usage much simpler, as after a few empirical runs the total memory needed is whatever spill memory is set, plus whatever is measured for non-spillable.

[ding-young]

@milenkovicm Thank you for sharing the experience.

The other reason is that we can’t do much with them; from our experience, they freak out a bit too early before a spill can occur.

Yes, when a non-spillable operator runs out of memory, it’s difficult to trigger spilling in another spillable operator to reclaim memory which seems to be a limitation currently.

Have you ever observed any cases where a non-spillable operator showed a memory usage spike (for example, due to skewness or similar factors)?

I wonder what would be the solution for these frequent failures on non-spillable operators - especially when other concurrent operators are spillable. If the memory usage of non-spillable operators can be roughly estimated before execution, do you think it would make sense to bypass or pre-reserve memory for them, instead of continuously growing the shared memory reservation along the non-spillable path?

[milenkovicm]

Yes, when a non-spillable operator runs out of memory, it’s difficult to trigger spilling in another spillable operator to reclaim memory which seems to be a limitation currently.

We need collaborative spilling like in spark, but I don't think current API can support it.

Have you ever observed any cases where a non-spillable operator showed a memory usage spike (for example, due to skewness or similar factors)?

I cant really say, we just track maximum memory used by non-spillable operators

I wonder what would be the solution for these frequent failures on non-spillable operators - especially when other concurrent operators are spillable. If the memory usage of non-spillable operators can be roughly estimated before execution, do you think it would make sense to bypass or pre-reserve memory for them, instead of continuously growing the shared memory reservation along the non-spillable path?

Tuning spillable and unspillable for me is like an equation with two unknowns, very hard to get it right as there is linear relation between two numbers. So we need to assume one variable as constant or add an additional equation. Adding an additional pool, one for spillable and another for non-spillable, like, I believe, @2010YOUY01 mentioned as well, will break relation between spillable and unspilable memory, simplifying things a bit, one variable can be approximated with a constant making relation with single variable.

Spillable pool is limited to trigger spill, unspillable may or may not be. If OS-level memory enforcement is used (cgroups), then unlimited unspillable may make sense. If we want to limit unspillable memory, doing a few empirical tests can give us rough estimation of.

[milenkovicm]

i should have said in my first comment, that we do have two different counts, but we keep unspillable limited by max_u64 (unlimited)

[16pierre]

Thanks for #18014 @adriangb, this could definitely fix the issue I previously hit here.

Back then I ended up working around without RepartitionExec nodes, and trashed my original drafts. My work-around was to manually control the parallelism on the datasources and coalesce some UnionExec nodes to control parallelism etc.

We could potentially close this issue as no longer relevant on our end. There's some interesting wider design discussions regarding how to handle the spillable vs non-spillable pools here, I quite like @milenkovicm approach here, next time we hit spilling problems we'll give it a shot, in previous iterations we definitely spent quite a few dev cycles fine-tuning the memory settings (especially before multi-level merge was released).

[adriangb]

Having played around with spilling in DataFusion now I think @milenkovicm makes some very good points. The way non-spill able operators work right now is not great.

What I would propose is that non-spillable operators register their memory usage but don't error if they blow past the memory limit. That would put pressure on spoilable operators to start spilling without running into these sort of memory reservation deadlocks. I think this is what @milenkovicm means as well by "we do not limit memory for non-spillable operators [...] Actually, we do track the memory usage of them, but we let it grow unbounded.".

Also what @ding-young said:

Yes, when a non-spillable operator runs out of memory, it’s difficult to trigger spilling in another spillable operator to reclaim memory which seems to be a limitation currently.

I think if we could make those two changes DataFusion would be in a much better place w.r.t. spilling. There would still be a long tail of making as many operators spoilable as possible and making them as performant when they do spill as we can but it becomes less important.

[2010YOUY01]

What I would propose is that non-spillable operators register their memory usage but don't error if they blow past the memory limit. That would put pressure on spoilable operators to start spilling without running into these sort of memory reservation deadlocks. I think this is what @milenkovicm means as well by "we do not limit memory for non-spillable operators [...] Actually, we do track the memory usage of them, but we let it grow unbounded.".

Probably we can let non-spillables go beyond a soft memory limit, but fails under a hard limit. Here are the rationales:

Now non-spillable memory reservations are used in the operators that can use unbounded memory, but haven't implemented a spilling yet, and it's possible to implement spilling in the future.

Letting non-spillable operators always continue even if OOMed works for the cases that the non-spillable operators can finish within the configured memory pool size, however now they're failing due to the current FairSpillPool's problematic design under race conditions.

This approach does not work if those non-spillable operators grow memory way beyond the configuration. All physical memory will be consumed before failure. Probably not an issue if we only run an query instance inside a isolated environment like docker, but it removes the opportunity for multiple concurrent queries to be registered in the same memory pool, which would enable more efficient memory sharing.

An alternative I can think of is setting a soft memory limit with a hard memory limit (for example under 8GB soft memory limit, non-spillables are safe to go beyond it, and if it happens, spillables will have some response time to do the spill writes and free up memory; if they go beyond the hard memory limit -- let's say soft-limit * 120% -- the query fails.

Another approach is letting spillables and non-spillables execute cooperatively -- if non-spillable OOMed, it will wait for spillables to do the spill writes, and wake up retry. But I think it would be very hard to implement.