NoahKusaba opened a new pull request, #2803:
URL: https://github.com/apache/datafusion-ballista/pull/2803
# Which issue does this PR close?
No issue. Found while measuring the sort-shuffle writer's memory accounting
for #2774.
# Rationale for this change
At the end of input, the sort-shuffle writer encodes its buffered rows into
one
IPC buffer per output partition. The buffered batches stay alive until the
last
partition is gathered, so the writer holds the batches and the encoded output
at the same time. Only the batches were reserved in the memory pool, so the
pool could grant the encoded output's memory to other consumers.
Measured with a counting allocator, 200 output partitions and 2M rows with a
string column: the writer held 287 MiB at its peak but had reserved 205 MiB,
82 MiB short. After this change it reserves 263 MiB at its peak. Most of the
remaining gap is allocator size-class rounding, which no reservation counts.
# What changes are included in this PR?
`encode_buffered_partitions` reserves each partition's encoded bytes as soon
as
that partition is encoded:
- If the pool grants them, the partition stays in memory.
- If the pool refuses, that partition is kept anyway, because its rows are
already encoded. Every later partition's rows go to the spill files through
the existing `SpillManager::spill`. This counts as one memory-pool spill
event, with its own debug log line, and is timed in `spill_time`.
- Afterwards the reservation holds only the encoded bytes, until the
consolidated write consumes them.
So at most one partition's encoded bytes are held beyond what the pool
grants:
the partition being encoded, or the one the pool refused.
Encoding now runs before the spill writers are finished, so spilled rows go
to
files that are still open.
## Alternatives considered
- **Reserve an estimate before encoding.** The only estimate available before
encoding is the buffered size. With lz4, that is about 3× the encoded
bytes,
so the writer would spill whole buffers that would have fit.
- **Write the refused partition's encoded bytes to its spill file instead of
keeping it.** This removes the one-partition overrun, but needs a new
`SpillManager` API for appending an already-encoded stream, and needs
`spill()` to append to a spill file instead of recreating it. That is about
100 more lines, only to handle the moment the pool refuses.
# Are there any user-facing changes?
No API or configuration changes. Two behaviour changes:
- Under memory pressure, the end of input can now spill. Before, it never
did. `spill_count` and disk I/O can therefore go up for jobs that run short
of pool memory.
- After encoding, the reservation drops from the buffered size to the encoded
size, which is usually smaller with compression. Other inputs of the same
task can use that memory sooner.
# How was this tested?
- `encoding_spills_the_partitions_after_the_pool_refuses`: encodes 4
partitions with room for partition 0 only. It checks which partitions stay
in memory and which spill, each partition's encoded bytes, and the bytes
left reserved. An unlimited pool keeps every partition in memory.
- `spills_when_the_pool_cannot_hold_the_encoded_output`: runs the writer with
a pool that holds the buffered rows but not the encoded output. It checks
that the writer spills once and that every partition's batches equal an
unlimited-pool run's.
- Removing the pool check, the final `resize`, or the switch to spilling each
makes at least one test fail.
- `cargo test -p ballista-core --lib`, `cargo clippy -p ballista-core
--all-targets -- -D warnings` and `cargo fmt --all -- --check`.
- This branch also builds and passes all tests when merged with #2800.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
