Repository navigation
Conversation
The output coalescer in RepartitionExec copies incoming rows into its own buffers until it has a full batch. Only completed batches were charged to the output partition's memory reservation, at send. The buffered rows were outside every reservation, so a memory pool limit did not see them. Charge the coalescer's size to the output partition's reservation after each push, once completed batches have been drained, and release it when the coalescer is finalized. A completed batch is still charged at send, after the coalescer's charge has dropped by the bytes it took, so each byte is charged once. Signed-off-by: Panagiotis Moustafellos <2493339+pmoust@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
When
RepartitionExeccoalesces its output, each output partition's coalescer copies incoming rows into its own buffers until it has a full batch. Only completed batches are charged to the output partition'sMemoryReservation, at send. The rows the coalescer holds before that are charged nowhere, so the memory pool does not see them. The amount grows with the number of output partitions and withbatch_size, and string columns make it larger, because the coalescer copies string view data into new buffers.What changes are included in this PR?
LimitedBatchCoalescer::size()passes through arrow'sBatchCoalescer::size().SharedCoalescerkeeps, next to the coalescer and under the same mutex, the number of bytes it has charged to the output partition's reservation.Completed batches are still charged in
OutputChannel::sendas before. The coalescer's charge drops by the bytes a completed batch took beforesendcharges that batch, so each byte is charged once.What is the testing strategy for this PR?
New test
repartition_reserves_coalescer_buffered_rowsinrepartition/mod.rs. Withbatch_size = 100, it pushes 150 rows as five small batches through aRepartitionExecand holds the input open after the last one. The first 100 rows form a completed batch; the last 50 stay in the coalescer. The test:LimitedBatchCoalescerfed the same batches reports as itssize(). This shows the buffered rows are charged and the completed batch is not counted a second time.Without the change, the first check fails: the pool reports 0 bytes reserved while the coalescer holds the 50 rows.
Are there any user-facing changes?
No API changes. Memory pools now see more of the memory that
RepartitionExecuses. Under a tight memory limit, repartition spills completed batches a little sooner, and other consumers see less free memory, because the buffered rows are now counted. A query that only just fit under its memory limit before can now fail withResourcesExhaustedin an operator that cannot spill.Open questions
grow, nottry_grow. The buffered rows cannot be spilled, so there is nothing to do when the pool refuses. As a result, the reservation can go over the pool limit by up to one coalescer's size per output partition (aboutbatch_sizerows each). Other consumers then see the pool as full and spill or fail. An alternative is to flush the partial batch early whentry_growfails and send it through the existing spill path. That puts a hard limit on the memory, but it produces smaller batches under memory pressure. It also needs a way to flushLimitedBatchCoalescerwithout finishing it. I kept the smaller change. I am happy to switch if reviewers prefer the hard limit.BatchCoalescer::size()counts the capacity of the in-progress buffers, and for primitive columns the coalescer reserves a fullbatch_sizeon the first write. So a partition holding a few rows is charged for a full batch's capacity. I think that is right, because that memory is allocated, but it is more than the bytes of the rows themselves.LIMIT), the input task drops that output channel without finalizing the coalescer. The charge then stays until the reservation is dropped with the last input task. The buffered memory has the same lifetime, so the accounting stays correct, but nothing releases it earlier.MemoryReservationupdates are atomic and no code path takes the coalescer mutex from inside a pool call, so I do not expect a lock-order problem. A customMemoryPoolwhosegrowblocks would now block input tasks of that output partition while they hold the mutex.