Indeterministic test failure in `SpillPool`

[2010YOUY01]

Describe the bug

The failed test:

datafusion/datafusion/physical-plan/src/spill/spill_pool.rs

Line 414 in af77197

/// assert_eq!(batches_read, 5);

It failed in one CI test run, however it can't be reproduced through repeated runs. It can be a heisenbug.
https://github.com/apache/datafusion/actions/runs/21394691626/job/61590390139?pr=20026

To Reproduce

No response

Expected behavior

No response

Additional context

No response

[dekuu5]

hey i am new to datafusion can i take my luck on this one ?

[2010YOUY01]

Thank you. The first step would be try to reproduce the test failure locally (e.g. run this test 100 times), and probably it requires some localized change to fix.

[dekuu5]

Hello @2010YOUY01,

I spent some time investigating this issue. Initially, I wasn't able to reproduce the bug even when running the tests 200 times in parallel.

However, I wrote a custom stress-test script to run the test case with much higher concurrency (100 parallel instances), and I was finally able to reproduce the failure consistently.

After debugging the reproduction, I identified a race condition in the coordination logic of the SpillPool. The poll_next function relies on a buffered stream (spawn_buffered) to read the file concurrently. The issue is that the background buffer task is not aware of the writer's status. Under heavy load, the buffer task can hit a temporary EOF (before the writer finishes) and quit prematurely. As a result, poll_next receives None from the stream and closes the reader before all batches are processed.
i changed the stream to be a normal unbuffered stream and it worked i will open a pr shortly.