From 6b9e91f0f0d4dba41399e2da0fa99f9ae62fd02d Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:08:32 +0000 Subject: [PATCH 1/3] Ensure stream is passed in Sort/SortBy exprs --- .../cudf_polars/dsl/expressions/sorting.py | 10 +++++--- python/cudf_polars/tests/test_groupby.py | 23 ++----------------- 2 files changed, 9 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/sorting.py b/python/cudf_polars/cudf_polars/dsl/expressions/sorting.py index b657f4522bb1..a3ca362219ba 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/sorting.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/sorting.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # TODO: remove need for this # ruff: noqa: D101 @@ -43,7 +43,7 @@ def do_evaluate( [descending], nulls_last=[nulls_last], num_keys=1 ) do_sort = plc.sorting.stable_sort if stable else plc.sorting.sort - table = do_sort(plc.Table([column.obj]), order, null_order) + table = do_sort(plc.Table([column.obj]), order, null_order, stream=df.stream) return Column( table.columns()[0], is_sorted=plc.types.Sorted.YES, @@ -80,6 +80,10 @@ def do_evaluate( ) do_sort = plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key table = do_sort( - plc.Table([column.obj]), plc.Table([c.obj for c in by]), order, null_order + plc.Table([column.obj]), + plc.Table([c.obj for c in by]), + order, + null_order, + stream=df.stream, ) return Column(table.columns()[0], dtype=self.dtype) diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index 6c1f328c6047..8eb39945873e 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -157,15 +157,7 @@ def test_groupby_sorted_keys( df: pl.LazyFrame, keys, exprs, - request, ): - request.applymarker( - pytest.mark.xfail( - is_streaming_engine(engine), - strict=False, - reason="https://github.com/NVIDIA/cudf/issues/21642 - no deterministic sort for keys", - ) - ) sorted_keys = [ key.sort(descending=descending) for key, descending in zip(keys, itertools.cycle([False, True])) @@ -173,20 +165,9 @@ def test_groupby_sorted_keys( q = df.group_by(*sorted_keys).agg(*exprs) - schema = q.collect_schema() - sort_keys = list(schema.keys())[: len(keys)] - # Multiple keys don't do sorting + sort_keys = list(q.collect_schema().keys())[: len(keys)] qsorted = q.sort(*sort_keys) - if len(keys) > 1: - # https://github.com/pola-rs/polars/issues/17556 - # Can't assert that the query without post-sorting fails, - # since it _might_ pass. - assert_gpu_result_equal(qsorted, engine=engine, check_exact=False) - elif schema[sort_keys[0]] == pl.Boolean(): - # Boolean keys don't do sorting, so we get random order - assert_gpu_result_equal(qsorted, engine=engine, check_exact=False) - else: - assert_gpu_result_equal(q, engine=engine, check_exact=False) + assert_gpu_result_equal(qsorted, engine=engine, check_exact=False) def test_groupby_len(engine: pl.GPUEngine, df, keys): From 3be2094294ee4f84bca31b868fb51b28db417451 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:23:51 +0000 Subject: [PATCH 2/3] Add more missing streams to pylibcudf calls --- .../cudf_polars/dsl/expressions/rolling.py | 15 +++++++++------ .../cudf_polars/dsl/expressions/unary.py | 2 +- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index 930683a3074f..72b5cb5dd32e 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -506,9 +506,9 @@ def _( # order_by expressions require us order each group lg = op.local_grouper assert isinstance(lg, plc.groupby.GroupBy) - _, rank_tables = lg.scan(rank_requests) + _, rank_tables = lg.scan(rank_requests, stream=df.stream) else: - _, rank_tables = grouper.scan(rank_requests) + _, rank_tables = grouper.scan(rank_requests, stream=df.stream) return rank_out_names, rank_out_dtypes, rank_tables @_apply_unary_op.register @@ -538,6 +538,7 @@ def _( # type: ignore[no-untyped-def] _, filled_tbl = local_grouper.replace_nulls( vals_tbl, [op.policy] * len(plc_cols), + stream=df.stream, ) tables = [plc.Table([column]) for column in filled_tbl.columns()] @@ -598,14 +599,14 @@ def _( # type: ignore[no-untyped-def] local_grouper = op.local_grouper assert isinstance(local_grouper, plc.groupby.GroupBy) - _, tables = local_grouper.scan(requests) + _, tables = local_grouper.scan(requests, stream=df.stream) result_tables: list[plc.Table] = [] for tbl, policy in zip(tables, fill_policies, strict=True): if policy is None: result_tables.append(tbl) else: - _, filled = local_grouper.replace_nulls(tbl, [policy]) + _, filled = local_grouper.replace_nulls(tbl, [policy], stream=df.stream) result_tables.append(filled) return out_names, out_dtypes, result_tables @@ -1275,7 +1276,7 @@ def do_evaluate( # noqa: D102 other_scalars, df, by_cols=by_cols ) - group_keys_tbl, value_tables = grouper.aggregate(gb_requests) + group_keys_tbl, value_tables = grouper.aggregate(gb_requests, stream=df.stream) broadcasted_cols = self._broadcast_agg_results( by_tbl, group_keys_tbl, @@ -1308,7 +1309,9 @@ def do_evaluate( # noqa: D102 order_sensitive, df, order_index=order_index, by_cols=by_cols ) - group_keys_tbl_local, value_tables_local = local.aggregate(gb_requests) + group_keys_tbl_local, value_tables_local = local.aggregate( + gb_requests, stream=df.stream + ) broadcasted_cols.extend( self._broadcast_agg_results( by_tbl, diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py index 2a6920f8c0e1..f899cb0a8e93 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py @@ -1197,7 +1197,7 @@ def do_evaluate( ] (keys_table, (counts_table,)) = plc.groupby.GroupBy( df.table, null_handling=plc.types.NullPolicy.INCLUDE - ).aggregate(gb_requests) + ).aggregate(gb_requests, stream=df.stream) if sort: sort_indices = plc.sorting.stable_sorted_order( counts_table, From b181c92924f74b396263b3878c96880e211267be Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Wed, 9 Sep 2026 04:08:40 +0000 Subject: [PATCH 3/3] Remove strict=False for test_non_coalesce_join --- python/cudf_polars/tests/test_join.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/python/cudf_polars/tests/test_join.py b/python/cudf_polars/tests/test_join.py index aa8ab27ee434..b71644ff5aab 100644 --- a/python/cudf_polars/tests/test_join.py +++ b/python/cudf_polars/tests/test_join.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -15,7 +15,6 @@ from cudf_polars.dsl import expr as ir_expr from cudf_polars.dsl.ir import ConditionalJoin from cudf_polars.testing.asserts import assert_gpu_result_equal -from cudf_polars.testing.engine_utils import is_streaming_engine @pytest.fixture(params=[False, True], ids=["nulls_not_equal", "nulls_equal"]) @@ -80,15 +79,7 @@ def test_non_coalesce_join( how, nulls_equal, join_expr, - request, ): - request.applymarker( - pytest.mark.xfail( - is_streaming_engine(engine), - strict=False, - reason="Non deterministic sort/join on nulls", - ) - ) query = left.join( right, on=join_expr, how=how, nulls_equal=nulls_equal, coalesce=False )