[BUG] top_k incorrect with non-elementwise `by` expression for cudf_polars + rapidsmpf
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
```python
In [1]: from functools import partialmethod
...: import polars
...: from cudf_polars.utils.config import Runtime, StreamingFallbackMode
...: executor = "streaming"
...: executor_options: dict[str, Any] = {}
...: executor_options["max_rows_per_partition"] = 4
...: executor_options["target_partition_size"] = 10
...: # We expect many tests to fall back, so silence the warnings
...: executor_options["fallback_mode"] = StreamingFallbackMode.SILENT
...: executor_options["runtime"] = Runtime.RAPIDSMPF
...: collect = polars.LazyFrame.collect
...: engine = polars.GPUEngine(executor=executor, executor_options=executor_options)
...: polars.LazyFrame.collect = partialmethod(collect, engine=engine)
...:
...: import polars as pl
...: pl.LazyFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8]}).top_k(
...: 2, by=(pl.when(pl.len() == 8).then(pl.col.a).otherwise(-pl.col.a))
...: ).collect()
Out[1]:
shape: (2, 1)
┌─────┐
│ a │
│ --- │
│ i64 │
╞═════╡
│ 1 │
│ 2 │
└─────┘
```
should be
```python
In [1]:
...: import polars as pl
...: pl.LazyFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8]}).top_k(
...: 2, by=(pl.when(pl.len() == 8).then(pl.col.a).otherwise(-pl.col.a))
...: ).collect()
Out[1]:
shape: (2, 1)
┌─────┐
│ a │
│ --- │
│ i64 │
╞═════╡
│ 8 │
│ 7 │
└─────┘
```
Contributor guide
Assessment
This issue has not been assessed yet.