apache / apache/datafusion

TopKRow ordering is inconsistent, preventing a possible sort optimization

Open
#24,659 1 comment 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.