TopKRow ordering is inconsistent, preventing a possible sort optimization
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
There is already a note about this [in the PartialOrd impl](https://github.com/apache/arrow-datafusion/blob/97d8c0d905b82ac3228453cbff185ab6ad646ef0/datafusion/physical-plan/src/topk/mod.rs#L1161). This is not consistent with PartialEq, which compares `row`, `batch_id` and `index` (that feels weird on its own; you'd think `row` is enough for eq in this context).
### To Reproduce
Attempting to optimize heap drains can change ordering, which makes some sqllogictest(s) fail since duplicate rows get drained in a new order. I ran into this when looking into TopK performance. Turns out that `BinaryHeap::into_sorted_vec` is slow (see https://github.com/rust-lang/rust/issues/115357) and doesn't have any particular ordering guarantees when draining.
See the test failures on [this CI run](https://github.com/apache/datafusion/actions/runs/32773884039/job/97582194027?pr=24639).
### Expected behavior
Ord impl on TopKRow accounts for `batch_id` and `index` to stabilize ordering of equal elements.
Current tests won't pass with this change though, making it a breaking change.
Here are test failures with stable ordering: [fail2.txt](https://github.com/user-attachments/files/31417120/fail2.txt)
### Additional context
In our case `let vec = heap.into_vec(); vec.sort_unstable(); vec` is around 55% faster than `heap.into_sorted_vec()` for 50000 elements (courtesy of [a new bench](https://github.com/massive-com/arrow-datafusion/blob/97d8c0d905b82ac3228453cbff185ab6ad646ef0/datafusion/core/benches/topk_aggregate.rs#L521-L526)). Stabilized ordering would unlock that optimization; `BinaryHeap::into_sorted_vec` and `Vec::sort_unstable` produce the same order with stable ordering. Structures other than BinaryHeap could also be used in the future. I've been looking into min-max heaps but could only get a 3-5% improvement over the BinaryHeap in std.
Another solution could be making `sort_unstable` opt-in with a config, which would avoid the breaking change. Unless we can rely the fact that TopK is already opt-in?
I think accounting for `batch_id` and `index` could be good future proofing though (see expected behavior above).
Contributor guide
Research direction
Start in datafusion/physical-plan/src/topk/mod.rs at the PartialOrd and PartialEq implementations for TopKRow, then review the TopK aggregate benchmark in datafusion/core/benches/topk_aggregate.rs. Confirm how heap draining and duplicate rows affect ordering, and run the relevant sqllogictests and benchmark to verify stable ordering and the proposed optimization.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100