Repository navigation
Memory account not adding up in SortExec #10073
Description
Describe the bug
- This is related / tangential to Add spilling in SortMergeJoin #9359
My primary problem is that I am trying to sort 100 million strings and always getting errors like:
Resources exhausted: Failed to allocate additional 39698712 bytes for ExternalSorterMerge[0] with 4053021824 bytes already allocated - maximum available is 26915520
After reading through the sort impl a bit I have noticed a few concerns (recorded in additional context)
To Reproduce
I don't have a df reproduction but this reproduces it for me in lance:
import pyarrow as pa
import lance
# Once the dataset has been generated you can comment out this generation code and reproduce the issue very quickly
print("Generating data")
my_strings = [f"string-{i}" * 3 for i in range(100 * 1024 * 1024)]
my_table = pa.table({"my_strings": my_strings})
print("Writing dataset")
ds = lance.write_dataset(
my_table, "/tmp/big_strings.lance", mode="overwrite", schema=my_table.schema
)
del my_table
# End of generation code
ds = lance.dataset("/tmp/big_strings.lance")
print("Training scalar index")
# To create a scalar index we must sort the column, this is where the error occurs
ds.create_scalar_index("my_strings", "BTREE")
Expected behavior
I can sort any number of strings, as long as I don't overflow the disk
Additional context
Here is how I understand memory accounting in the sort today:
- As batches arrive, they are placed in an accumulation queue, and the size of the spillable reservation grows
- Once the pool is out of space we begin the spill process
- The first part of the spill process is to sort the accumulation queue (which, at this point, has many batches in it)
- Each batch becomes an input stream for a SortPreservingMergeStream (this is a LOT of inputs).
- When the batch input stream is polled the batch is sorted, the unsorted batch is dropped (and removed from the spillable reservation) and the sorted batch is returned. The SortPreservingMergeStream then puts this batch in the batch builder (which adds it to the non-spillable reservation). This is a problem, as described below
- As the sort preserving merge stream is polled it polls the various inputs, fills up the batch builder, and then starts to emit output batches
- Back in the sort exec the sort preserving merge stream is fully drained (try_collect). This is a problem, as described below
- These collected and sorted batches are then written to the spill file.
The first problem (and the one causing my error) is that a sorted batch of strings (the output of sort_batch) is occupying 25% more memory than the unsorted batch of strings. I'm not sure if this buffer alignment, padding, or some kind of 2x allocation strategy used by the sort, but it seems reasonable something like this could happen. Unfortunately, this is a problem. We are spilling because we have used up the entire memory pool. We now take X bytes from the memory pool, convert it into 1.25 * X bytes, and try to put it back in the memory pool. This fails with the error listed above.
The second problem is that we are not accounting for the output of the sort perserving merge stream. Each output batch from the sort preserving merge stream is made up of rows from the various input batches. In the degenerate case, where the input data is fully random, this means we will probably require 2 * X bytes. This is because each output batch is made up of 1 batch from each input stream. We can't release any of the input batches until we emit the final output batch.
The solution to this second problem is that we should be streaming into the spill file. We should not collect from the sort preserving merge stream and then write the collected batches into the spill file. This problem is a bit less concerning for me at the moment because it is "datafusion uses more memory than it should" and not "datafusion is failing the plan with an error". We don't do a lot of sorting in lance and so we can work around it reasonably well by halving the size of the spill pool.
Activity
alamb commentedon Apr 15, 2024
It sounds like the issue, at a (really) high level is "additional buffer space is required to actually implement the spill"
And since during spill the plan is under memory pressure, getting this additional memory can and does fail
Some strategies I can think of are:
Maybe we can do 1 in the sort term while figuring out a more sophisticated strategy for 2 or 3
comphead commentedon Apr 15, 2024
I'm investigating this as part of #9359
yjshen commentedon Aug 25, 2024
I find this also related. #9528 (comment)
yjshen commentedon Aug 25, 2024
Another point of code worth noticing is inside the current
sort_batchimplementation:datafusion/datafusion/physical-plan/src/sorts/sort.rs
Lines 601 to 607 in 79fa6f9
Performance-wise, I think it's beneficial to apply the row format comparison to all multi-column cases, however, while considering
sort_batchis used in multiple places wherespillis being called; creating rows before comparing would introduce more memory pressure.BTW, I think we should report memory usage inside
lexsort_to_indices_multi_columns:datafusion/datafusion/physical-plan/src/sorts/sort.rs
Lines 650 to 652 in 79fa6f9
alamb commentedon Nov 21, 2024
I think @2010YOUY01 may have fixed this recently 🤔
bfcrampton commentedon Jan 6, 2025
FWIW I'm still seeing the same issue through LanceDB (lance-format/lance#2119 (comment)).
westonpace commentedon Jan 6, 2025
This isn't necessarily indicative as Lance lags behind Datafusion (currently we are at 42 which is 4 months behind). However, I just updated my local lance to release 44 (which should contain the potential fix @alamb is alluding to) and confirmed that the issue is still not fixed.
This also doesn't surprise me. I think the issue here is not double-counting but rather is dealing with the fact that a string array uses more memory after sorting than it was using before sorting (and so we run out of memory trying to spill).
I'll try and find some time today to create a pure datafusion reproducer.
westonpace commentedon Jan 6, 2025
Here's a pure-rust datafusion-only example: westonpace@26ed75c
It takes a bit of time the first run to generate the strings test file (it probably doesn't need to be so big). After that it reproduces the issue quickly.
I've also added some prints that hopefully highlight the issue. Before we do an in-memory sort we have ~5MB of unsorted string data. After sorting we have 8MB of sorted string data.
This is not surprising to me. During the sort we are probably building a string array and probably using some kind of resize-on-append string building that is doubling and we end up with ~8MB because the amount we need is between 4MB and 8MB.
Unfortunately, this leads to a failure which it probably should not do. I think @alamb had some good suggestions in this comment
alamb commentedon Jan 7, 2025
This is a good call - maybe there is a factor of 2 less memory if we allocated the correct capacity up front somehow.
alamb commentedon Feb 15, 2025
@Kontinuation has a PR with a proposed fix: