Fix slicing after Join and GroupBy in streaming cudf-polars

[rjzamora]

Description

The streaming executor does not properly handle slicing after Join or GroupBy. Rather than slicing the "reduced" join/gropuby result, each partition is sliced individually. This PR includes a general fix, by pulling the slice operation out of Join and GroupBy nodes.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[TomAugspurger]

edit: probably not an issue. See #19187 (comment)

@rjzamora there's one similar issue on branch-25.08 that doesn't appear to be fixed here:

import polars as pl

from cudf_polars.testing.asserts import assert_gpu_result_equal


left = pl.LazyFrame(
        {
            "a": [1, 2, 3, 1, None],
            "b": [1, 2, 3, 4, 5],
            "c": [2, 3, 4, 5, 6],
        }
    )
right = pl.LazyFrame(
        {
            "a": [1, 4, 3, 7, None, None, 1],
            "c": [2, 3, 4, 5, 6, 7, 8],
            "d": [6, None, 7, 8, -1, 2, 4],
        }
    )

q = left.join(right, on="a", how="left", nulls_equal=False, coalesce=True).slice(1, 5)
assert_gpu_result_equal(q, engine=pl.GPUEngine(executor="streaming", executor_options={"max_rows_per_partition": 4}))

which fails with

AssertionError: DataFrames are different (value mismatch for column 'a')
[left]:  [1, 2, 3, 1, 1]
[right]: [1, 1, 2, None, 1]

I'm starting to investigate that now. Should we handle that in a separate PR?

[TomAugspurger]

Maybe this is just a join ordering issue. I don't think that left.join(right).slice(...) makes sense unless you've specified a maintain_order. Otherwise, you're just slicing random rows, as chosen by the join implementation.

We do match for this snippet that specifies maintain_order (ignore this; we only match because we don't implement maintain_order and fall back to CPU)

import polars as pl
from cudf_polars.testing.asserts import assert_gpu_result_equal

left = pl.LazyFrame(
    {
        "a": [1, 2, 3, 1, None],
        "b": [1, 2, 3, 4, 5],
        "c": [2, 3, 4, 5, 6],
    }
)
right = pl.LazyFrame(
    {
        "a": [1, 4, 3, 7, None, None, 1],
        "c": [2, 3, 4, 5, 6, 7, 8],
        "d": [6, None, 7, 8, -1, 2, 4],
    }
)

engine = pl.GPUEngine(
    executor="streaming", executor_options={"max_rows_per_partition": 4}
)
q = left.join(right, on="a", how="left", maintain_order="left").slice(1, 5)
assert_gpu_result_equal(q, engine=engine)

so I'll just plan to update #19146 to adjust these tests somehow, probably by specifying the engine explicitly...

[rjzamora]

Maybe this is just a join ordering issue. I don't think that left.join(right).slice(...) makes sense unless you've specified a maintain_order. Otherwise, you're just slicing random rows, as chosen by the join implementation.

Yes, this is exactly how I understand it.

[TomAugspurger]

Sorry, #19186 caused some failures here. I think just removing the config_options should suffice.

[TomAugspurger]

/merge