Repository navigation
[BUG] over(...) inside .filter() fallback reorders rows on multi-rank streaming engines #22405
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
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:
k: [x, y, x, y]
v: [3, 2, 5, 7]
Actual
The streaming engine returns the same rows, but reordered:
k: [x, x, y, y]
v: [3, 5, 2, 7]
Each RankActor also emits the fallback warning:
UserWarning: over(...) inside filter is not supported for multiple partitions;
falling back to in-memory evaluation.
Workaround
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.
Activity
wence- commentedon May 19, 2026
This is really an issue with
concat, not filter. See #22376Add regression test for NVIDIA#22405 order after concat
aryansri05 commentedon Oct 19, 2026
Following up on @wence-'s point that this is really the
concatordering bug: #24685 fixes #22376 by making an ordered multi-rankUnionpreserve global row order. With that PR the reproducer above passes onRayEngine(num_ranks=2)(onmainit fails with a value mismatch on columnk), and I addedtest_over_in_filter_after_concat_preserves_orderso it's covered on all engines.The existing xfails that point at this issue fail for other reasons, so I left them in place:
test_over_in_filter_unsupportedandtest_over_noncol_key_fallback(test_rolling.py): the fallback warning is raised on the Dask/Ray workers, sopytest.warnsnever sees it.test_join_and_slice(test_join.py): flaky onmaintoo, thezlice1cases pass or fail from run to run.Should those move to their own issue(s)? Happy to open them and update the xfail reasons, and this one could then be closed once #24685 lands.
Address review feedback on ordered Union