Busy-waiting in SortPreservingMergeStream

[pepijnve]

Describe the bug

When running a query like select a from annotated_data_infinite2 order by b desc limit 10, a SortPreservingMergeStream is created that merge sorts the presorted partitions. During the initial sorting phase of the query, the query execution will continuously hit

datafusion/datafusion/physical-plan/src/sorts/merge.rs

Line 236 in 1daa5ed

return Poll::Pending;

where the stream wakes the current task for rescheduling. Since all the input streams will be pending for quite some time during the sort of a large data the execution ends up in a busy-waiting loop.

I believe this was introduced in f2a8b07. It's not clear to me why the self wake was added.

To Reproduce

• Add print!("P"); stdout().flush(); just before the return at https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/sorts/merge.rs#L236
• Run a sort on a large dataset with target partitions > 1 so that SortPreservingMergeStream is used.
• Observe how often the poll_next_inner is hit

Expected behavior

• No busy-waiting during query execution

Additional context

No response

[pepijnve]

Investigating, but help appreciated.

[pepijnve]

A closer look shows I might be completely mistaken. Will close if irrelevant. Nope, not a case of PEBKAC after all. I've refined the reproduction scenario.

[pepijnve]

Just to have a hyperlink between the issues: as far as I can tell the behavior was introduced by #12302 to resolve #12300.