Open beckernick opened 3 months ago
Since (for now) there are no partial_sort/quickselect
algorithms in CUB/thrust (see https://github.com/NVIDIA/cccl/issues/931 and https://github.com/NVIDIA/cccl/issues/685), this needs to be implemented via full sort + slice.
This is best done by (rather than implementing top_k
as a special-case in the UnaryFunction
handler) translating a call to top_k
into Sort + slice
(note we need to implement slice
for that too, but that's relatively straightforward).
Polars actually doesn't guarantee that the returned values are sorted, just that they are the top_k, so that gives freedom in the implementation if/when something is implemented in CCCL
The top_k expression is essentially an nlargest + limit operation and is sometimes used to short circuit sort + limit scenarios common in SQL. E.g.,
SELECT * FROM table ORDER BY col DESC LIMIT 10
and the equivalent in DataFrame operations.(Related but separate from #16222)