Skip to content

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

Open

Description

@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:

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.

Activity

  1. 2010YOUY01 commented on Aug 28, 2025

    @2010YOUY01
    Member

    The current FairSpillPool implementation seems problematic
    See its behavior:

    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.

  2. 16pierre commented on Aug 28, 2025

    @16pierre
    Author

    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.

  3. 16pierre commented on Aug 28, 2025

    @16pierre
    Author

    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.

  4. ding-young commented on Aug 29, 2025

    @ding-young
    Contributor

    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.

  5. 16pierre commented on Aug 29, 2025

    @16pierre
    Author

    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:

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

    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:

    //! 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.

  6. milenkovicm commented on Oct 11, 2025

    @milenkovicm
    Contributor

    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.

  7. ding-young commented on Oct 11, 2025

    @ding-young
    Contributor

    @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?

  8. milenkovicm commented on Oct 11, 2025

    @milenkovicm
    Contributor

    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.

  9. milenkovicm commented on Oct 11, 2025

    @milenkovicm
    Contributor

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

  10. 16pierre commented on Oct 23, 2025

    @16pierre
    Author

    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).

  11. adriangb commented on Oct 25, 2025

    @adriangb
    Contributor

    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.

  12. 2010YOUY01 commented on Oct 26, 2025

    @2010YOUY01
    Member

    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.

  13. milenkovicm commented on Oct 26, 2025

    @milenkovicm
    Contributor

    Would it make sense to consider soft limit as spillble limit and hard limit spillable and non spillable combined?

    This way we break relation between two numbers, make it easier to reason about.

    We can make hard limit unlimited, it makes easier to estimate total number in tests or bound it in production if we wish

    Edit: getting over soft limit in total spill and unspill means that we should expect spilling to be enforced at some point not that it's enforced

  14. 12 remaining items

  15. milenkovicm commented on Nov 5, 2025

    @milenkovicm
    Contributor

    Just one illustration of a problem

    CustomSpillPool [ spillable (current/max/limit) 0.0 B / 1727.7 MB / 2.8 GB, un-spillable  (current/max/limit) 0.0 B / 231. 5MB / 16777216.0TB, spill count: 236]
    

    This is FairSpill (kind of), there are 236 failed allocations (spills) unspillable memory was never close to 2.8G which was the limit.

    Also note, total memory for process (writing to a local file) was up to 5GB (with Mimalloc allocator)

  16. milenkovicm commented on Nov 6, 2025

    @milenkovicm
    Contributor

    Maybe the simplest change to mimic some kind of cooperative spilling is to have a global state in MemoryPool, like a counter or timestamp, which would indicate when was the last time a consumer failed to grow allocation, lets call it "failed allocation counter".

    MemorConsumer apart from providing can_spill should provide last_spill_counter so if the counter is less than memory pool counter memory pool could fail to grow memory reservation and force consumer to spill. Implementations could make this logic customized, like spill all or spill until we get under 50% of reserved memory, or whatever ...

    This could be backward compatible change and we could get some kind of cooperative spilling without big impact, having a mutex state "failed allocation counter" and MemorConsumer::advance_counter which would be triggered by MemoryPool in case of failed allocation (instead of relying on memory consumer to advance it when it spills)

    I also believe we would need some kind of soft and hard memory, where unspillable consumers should get memory allocation if we're over soft limit but can trigger spill with incrementing "failed allocation counter"

    Lastly we would probably need some kind of grace period between "failed allocation counter" increment so it does not trigger too many unnecessary spills (not sure about this part)

  17. Kontinuation commented on Nov 24, 2025

    @Kontinuation
    Member

    I'm also looking forward to the cooperative spilling feature. This is important for projects such as Comet to implement Photon-like memory management and reduce the possibility of allocation failure (see related issue apache/datafusion-comet#949). A similar DataFusion-based Spark accelerator project Apache Auron re-implemented memory-intensive operators such as sort, aggregation, and join all by themselves to use their own memory manager and bypass the limitation of DataFusion's memory management API.

    A bit of history: the initial memory management proposal and implementation did support cooperative spilling. However, a later simplification removed that feature.

  18. adriangb commented on Nov 24, 2025

    @adriangb
    Contributor

    I was thinking about this a bit today. The thought I had was of some structure like (note this is a half baked idea):

    trait ExecutionPlan {
        async fn spill(&self, request_bytes: usize) -> Result<Option<usize>>> { Ok(None) }
        // existing stuff
    }
    
    struct CooperativeMemoryPool {
        plan: Arc<dyn ExecutionPlan>
        pool: Arc<dyn MemoryPool>
    }
    
    impl MemoryPool for CooperativeMemoryPool {
        async fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
            // Maybe some sort of re-entrant lock here?
            if let Ok(()) = self.pool.try_reserve(reservation, additional) {
                return Ok(())
            }
            // Do a top-down traversal through the plan asking each node to spill the required memory
           // until we've spilled enough or we realize we can't spill in which case error
           // Now try again?
        }
    }
  19. Kontinuation commented on Nov 24, 2025

    @Kontinuation
    Member

    I still prefer the MemoryConsumer abstraction in both Apache Spark and the initial MemoryConsumer in DataFusion. The memory pool coordinates between spillable data structures instead of execution plans. One execution plan instance could create multiple SendableRecordBatchStreams, and each stream may hold some instances of spillable data structures. We want the memory manager to coordinate the memory usage of these memory consumers directly instead of doing this through plans.

    The memory consumer must be able to differentiate between self-initiated spills (triggered by its own try_grow request) and passive spills (triggered cooperatively by other consumers), unless we ensure that the consumer calling try_grow will never have its spill method called in response to memory shortage.

    There are some common rules for implementing the MemoryConsumer interface in Apache Spark. The description of apache/spark#9241 has listed all the rules and considerations to avoid deadlocking, but I'd like to elaborate on them according to what I understand.

    1. Two stages of spillable operators

    Spillable operators are pipeline breakers (sort, grouped aggregation, etc.). These operators work in 2 stages:

    1. Build stage: ingest input batches and build some internal state. For the sort operator, the internal state consists of multiple sorted runs. Operators cannot produce batches until they have ingested all the input batches.
    2. Produce stage (the Spark PR calls it "scanning stage"): produces batches from their internal states. For instance, merge sorted runs as final ordered outputs.

    2. Cooperative spill only applies to the produce stage

    The operators usually intensively mutate the internal state in the build stage; it is no fun to handle concurrent passive spill requests in this stage. The operators only self-spill in build stage and ignore any spill requests initiated by other consumers.

    The internal state becomes mostly read-only in produce stage. Handling passive spill requests is much easier. Here is Spark's sorting implementation as an example: it checks whether the spill is triggered by itself (self-spill). If it is not, it handle the spill request only when readingIterator is available, which is an indication of being in produce stage.

    Here is an illustration of a fragment of a query plan involving multiple spillable operators, one is in produce stage and the other is in build stage:

             ...
              |
              v
    +-------------------+
    | GroupedAggregate  | in produce stage, can respond to spill request when Sort fails to reserve memory.
    +-------------------+
              |
              v
    +-------------------+
    |       Sort        | in build stage, can self-spill, or ask GroupedAggregate to spill.
    +-------------------+
              |
              v
             ...
    

    3. Don't allocate memory when spilling

    The documentation of MemoryConsumer.spill explicitly required that consumers should not allocate memory when spilling. This is for preventing deadlocks.

    4. Do not free memory when holding the lock of self

    We usually need to protect the internal state using mutexes to support handling concurrent cooperative spilling requests while producing batches. We also usually hold the lock in spill method to mutate the internal state to free some memory. One caveat is that we should only call memory management methods such as free or shrink after the lock is released. There is an example in Spark's sort operator: it deliberately postpones calls to freePage after the synchronize block.

  20. milenkovicm commented on Nov 24, 2025

    @milenkovicm
    Contributor

    I believe @Kontinuation’s point is spot on. Most memory issues I encounter stem from a spillable operator (A) transitioning to producing state while another operator (B) begins consuming.

  21. nathanb9 commented on Apr 17, 2026

    @nathanb9
    Contributor

    @milenkovicm @Kontinuation @2010YOUY01 @adriangb
    Sorry to ping you all! Im trying to get your thoughts on a interface im proposing to allow cooperative spilling. I think this could help address some things you've all discussed here or at least restart some of the discussion.

    Here is PR please take a look:#21425

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions