NVIDIA / NVIDIA/cudf

[BUG] `over(...)` inside `.filter()` fallback reorders rows on multi-rank streaming engines

Open
#22,405 1 comment 0 reactions 0 assignees View on GitHub
bug cudf-polars dask rapids-mpf
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

### Description

When a query uses `pl.col(...).over(...)` (or `pl.len().over(...)`) inside `.filter(...)` on a multi-rank streaming engine (`RayEngine`, `DaskEngine`, or `SPMDEngine` with `num_ranks > 1`), cudf-polars falls back to in-memory evaluation; see `cudf_polars/experimental/utils.py:88`.

The fallback executes independently on each rank, then concatenates the per-rank results. This preserves the correct set of rows, but not the original global row order from the input scan/concat.

As a result, `assert_gpu_result_equal(..., check_row_order=True)` (the default) fails with a row-order mismatch.

### Reproducer

```python
import polars as pl

from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine
from cudf_polars.testing.asserts import assert_gpu_result_equal

q = pl.concat(
[
pl.LazyFrame({"k": ["x", "y"], "v": [3, 2]}),
pl.LazyFrame({"k": ["x", "y"], "v": [5, 7]}),
]
).filter(pl.len().over("k") == 2)

with RayEngine(
num_ranks=2,
executor_options={"max_rows_per_partition": 1},
engine_options={"allow_gpu_sharing": True},
ray_init_options={"include_dashboard": False},
) as eng:
# Raises: DataFrames are different (value mismatch for column "k")
assert_gpu_result_equal(q, engine=eng)
```

### Expected

CPU and streaming execution should return the same four rows in the same order:

```text
k: [x, y, x, y]
v: [3, 2, 5, 7]
```

### Actual

The streaming engine returns the same rows, but reordered:

```text
k: [x, x, y, y]
v: [3, 5, 2, 7]
```

Each `RankActor` also emits the fallback warning:

```text
UserWarning: over(...) inside filter is not supported for multiple partitions;
falling back to in-memory evaluation.
```

### Workaround

```python
assert_gpu_result_equal(q, engine=eng, check_row_order=False)
```

The data is correct; only the row ordering differs.

### Possible fix

The per-rank fallback evaluates `over(...)` on local partitions only, so original global row order is lost after concatenation.

Possible fixes:

* Preserve original row positions, for example by adding a temporary `__row_idx__` column before fallback execution and restoring order after concat.
* Explicitly document that `over(...)` inside `.filter(...)` does not preserve row order on multi-rank streaming engines.

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.