NVIDIA / NVIDIA/cudf

[BUG] top_k incorrect with non-elementwise `by` expression for cudf_polars + rapidsmpf

Open
#22,074 1 comment 0 reactions 0 assignees View on GitHub
bug cudf-polars
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.