From 1144a2ade7b7251bcc698670f5ae0a09a7a51eb6 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Mon, 16 Mar 2026 20:14:14 +0000 Subject: [PATCH 1/2] Change more Rapidsmpf Shuffler.wait_on to Shuffler.wait --- python/cudf_polars/cudf_polars/experimental/shuffle.py | 2 +- python/cudf_polars/cudf_polars/experimental/sort.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/experimental/shuffle.py index 0cd70ea375ef..efb134abf7cd 100644 --- a/python/cudf_polars/cudf_polars/experimental/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/shuffle.py @@ -104,7 +104,7 @@ def extract_partition( context = get_worker_context() - shuffler.wait_on(partition_id) + shuffler.wait() column_names = options["column_names"] dtypes = options["dtypes"] return DataFrame.from_table( diff --git a/python/cudf_polars/cudf_polars/experimental/sort.py b/python/cudf_polars/cudf_polars/experimental/sort.py index 2f0f531a3fc9..98b2cdb0f610 100644 --- a/python/cudf_polars/cudf_polars/experimental/sort.py +++ b/python/cudf_polars/cudf_polars/experimental/sort.py @@ -370,7 +370,7 @@ def extract_partition( context = get_worker_context() - shuffler.wait_on(partition_id) + shuffler.wait() column_names = options["column_names"] column_dtypes = options["column_dtypes"] From 609cf99eba923682fb81bdec8513c7a1bbc27afc Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Mon, 16 Mar 2026 20:15:22 +0000 Subject: [PATCH 2/2] Pass df.stream to to_py call in shift --- python/cudf_polars/cudf_polars/dsl/expressions/unary.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py index 70fc223d965c..275376c132a6 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py @@ -522,9 +522,7 @@ def do_evaluate( offset = n_expr.value else: n_col = n_expr.evaluate(df, context=context) - offset_py = plc.copying.get_element( - n_col.obj, 0, stream=df.stream - ).to_py() + offset_py = n_col.obj_scalar(stream=df.stream).to_py(stream=df.stream) assert isinstance(offset_py, int) offset = offset_py if self.name == "shift":