andygrove commented on code in PR #2803:
URL:
https://github.com/apache/datafusion-ballista/pull/2803#discussion_r4526425565
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -122,10 +130,18 @@ fn encode_buffered_partitions(
// One writer open at a time: each holds its own compression context.
buffered.drain(|partition, batches| {
let (buf, stat) = write_ipc_stream(Vec::new(), schema, opts, batches)?;
- encoded[partition] = buf;
- stats[partition] = stat;
+ if reservation.try_grow(buf.capacity()).is_ok() {
Review Comment:
Nice that this reserves the real allocation rather than `len()`. One
thought: `buf` grows by doubling inside `write_ipc_stream`, so its capacity can
be close to twice the bytes written. That means a partition that would fit can
still get spilled. Would a `buf.shrink_to_fit()` before the `try_grow` be worth
it? It costs a copy, but the writer would then hold less memory for the rest of
the task and ask the pool for less. Happy either way.
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -2243,6 +2286,92 @@ mod tests {
.await
}
+ #[tokio::test]
+ async fn spills_when_the_pool_cannot_hold_the_encoded_output() ->
Result<()> {
Review Comment:
This test depends on the buffered rows coming to just under 193 KiB. If the
`BufferedBatches` accounting ever grows a little, the input loop spills
instead. The test would still pass but would no longer cover the encode path.
`encoding_spills_the_partitions_the_pool_refuses` covers that path directly, so
this is minor. Could the test also check that the spill came from encoding, for
example that the input loop recorded no spill event? Or a comment saying the
margin is tight would help the next person who changes the accounting.
--
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]
