NoahKusaba commented on code in PR #2803:
URL:
https://github.com/apache/datafusion-ballista/pull/2803#discussion_r4530188649
##########
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:
Left it out: under mimalloc, `realloc` keeps the same block when shrinking
to at least half its size, so `shrink_to_fit` frees nothing but lowers
`capacity()`, and we'd under-reserve. Reserving `capacity()` over-counts
slightly, the safe side. Added a comment.
##########
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:
Both spill paths write identical bytes, so the test now asserts its
precondition: the buffered bytes, computed with the writer's own accounting,
fit the pool. If the accounting grows, the test fails instead of silently
covering the input loop.
--
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]
