Run cudf_polars unit tests with RapidsMPF

.github/workflows/pr.yaml

@@ -480,17 +480,6 @@ jobs:
       matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
       build_type: pull-request
       script: "ci/test_wheel_cudf_polars.sh"
-  wheel-tests-cudf-polars-with-rapidsmpf:
-    needs: [wheel-build-cudf-polars, changed-files]
-    secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels && fromJSON(needs.changed-files.outputs.changed_file_groups).neither_cudf_nor_dask_cudf
-    with:
-      # This selects "ARCH=amd64 + the latest supported Python + CUDA" to minimize CI usage.
-      # (rapidsmpf compatibility already validated in rapidsmpf CI)
-      matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
-      build_type: pull-request
-      script: "ci/test_cudf_polars_with_rapidsmpf.sh"
   cudf-polars-polars-tests:
     needs: [wheel-build-cudf-polars, changed-files]
     secrets: inherit

.github/workflows/test.yaml

@@ -153,19 +153,6 @@ jobs:
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       script: "ci/test_wheel_cudf_polars.sh"
-  wheel-tests-cudf-polars-with-rapidsmpf:
-    secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main
-    with:
-      # This selects "ARCH=amd64 + the latest supported Python + CUDA" to minimize CI usage.
-      # (rapidsmpf compatibility already validated in rapidsmpf CI)
-      matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
-      build_type: ${{ inputs.build_type }}
-      branch: ${{ inputs.branch }}
-      date: ${{ inputs.date }}
-      sha: ${{ inputs.sha }}
-      script: "ci/test_cudf_polars_with_rapidsmpf.sh"
-      continue-on-error: true
   cudf-polars-polars-tests:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main

ci/run_cudf_polars_pytests.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
@@ -9,17 +9,32 @@ set -euo pipefail
 # Support invoking run_cudf_polars_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/
 
-# Test the "in-memory" executor
+echo "Test the in-memory executor"
 python -m pytest --cache-clear "$@" tests --executor in-memory
 
-# Test the default "streaming" executor
+echo "Test the default streaming executor"
 python -m pytest --cache-clear "$@" tests --executor streaming
 
-# Test the "streaming" executor with small blocksize
+echo "Test the streaming executor with small blocksize"
 python -m pytest --cache-clear "$@" tests --executor streaming --blocksize-mode small
 
-# Run experimental tests with Distributed cluster
+echo "Test the (future default) streaming executor with rapidsmpf"
+CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE=1 CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD=rapidsmpf python -m pytest --cache-clear "$@" tests \
+    --executor streaming \
+    --blocksize-mode small \
+    --cluster single \
+    --runtime rapidsmpf \
+    --cov-fail-under=0  # TODO? missing coverage may be due to old task-based paths
+
+echo "Run experimental tests with Distributed cluster"
 python -m pytest --cache-clear "$@" "tests/experimental" \
     --executor streaming \
     --cluster distributed \
     --cov-fail-under=0  # No code-coverage requirement for these tests.
+
+echo "Run experimental tests with the distributed cluster mode and the rapidsmpf runtime"
+python -m pytest --cache-clear "$@" "tests/experimental" \
+    --executor streaming \
+    --cluster distributed \
+    --runtime rapidsmpf \
+    --cov-fail-under=0  # No code-coverage requirement for these tests(?)

ci/test_wheel_cudf_polars.sh

@@ -30,7 +30,7 @@ rapids-pip-retry install \
     -v \
     --constraint ./constraints.txt \
     --constraint "${PIP_CONSTRAINT}" \
-    "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,experimental]" \
+    "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,experimental,rapidsmpf]" \
     "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
     "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)"
 

dependencies.yaml

@@ -115,6 +115,8 @@ files:
       - depends_on_cudf_kafka
       - depends_on_custreamz
       - depends_on_cudf_polars
+      - depends_on_rapidsmpf
+      - cudf_polars_trace
   test_java:
     output: none
     includes:

python/cudf_polars/cudf_polars/experimental/dask_registers.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """Dask function registrations such as serializers and dispatch implementations."""
@@ -215,7 +215,7 @@ def _(x: DataFrame) -> int:
         from rapidsmpf.integrations.dask.spilling import register_dask_serialize
 
         register_dask_serialize()  # pragma: no cover; rapidsmpf dependency not included yet
-    except ImportError:
+    except ImportError:  # pragma: no cover
         pass
 
     # Register the tokenizer for NamedExpr and DataType. This is a performance

python/cudf_polars/tests/conftest.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
@@ -28,6 +28,11 @@ def clear_memory_resource_cache():
     cudf_polars.callback.default_memory_resource.cache_clear()
 
 
+@pytest.fixture
+def using_rapidsmpf():
+    return cudf_polars.testing.asserts.DEFAULT_RUNTIME == "rapidsmpf"
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--executor",

python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py

@@ -20,6 +20,7 @@
 @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.")
 def test_structlog_streaming_node_events():
     """Test that structlog emits 'Streaming Actor' events when tracing is enabled."""
+    pytest.importorskip("structlog")
     # Run in subprocess to control CUDF_POLARS_LOG_TRACES environment variable
     code = textwrap.dedent("""\
     import polars as pl

python/cudf_polars/tests/expressions/test_agg.py

@@ -194,7 +194,13 @@ def test_implode_agg_unsupported():
     assert_ir_translation_raises(q, NotImplementedError)
 
 
-def test_decimal_aggs(decimal_df: pl.LazyFrame) -> None:
+def test_decimal_aggs(decimal_df: pl.LazyFrame, using_rapidsmpf, request) -> None:
+    request.applymarker(
+        pytest.mark.xfail(
+            condition=using_rapidsmpf,
+            reason="https://github.com/rapidsai/cudf/issues/21645",
+        )
+    )
     q = decimal_df.with_columns(
         sum=pl.col("a").sum(),
         min=pl.col("a").min(),

python/cudf_polars/tests/expressions/test_len.py

@@ -30,7 +30,13 @@ def test_len(dtype, empty):
 
 
 @pytest.mark.parametrize("data", [[1, 2, 3], [1, 2, None]])
-def test_col_len(data):
+def test_col_len(data, using_rapidsmpf, request):
+    request.applymarker(
+        pytest.mark.xfail(
+            condition=using_rapidsmpf,
+            reason="https://github.com/rapidsai/cudf/issues/21645",
+        )
+    )
     data = {"a": list("xyz"), "b": data}
     q = pl.LazyFrame(data).select(
         pl.col("a").len().alias("l"),

python/cudf_polars/tests/test_config.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import os
 import sys
 from typing import Any, cast
 
@@ -381,8 +382,12 @@ def test_scheduler_deprecated() -> None:
 
 def test_validate_shuffle_method_defaults(
     *,
+    monkeypatch: pytest.MonkeyPatch,
     rapidsmpf_distributed_available: bool,
 ) -> None:
+    # TODO: Can remove once removed from ci/run_cudf_polars_pytests.sh
+    # and the default shuffle method changes to "rapidsmpf"
+    monkeypatch.delenv("CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD", raising=False)
     config = ConfigOptions.from_polars_engine(
         pl.GPUEngine(
             executor="streaming",
@@ -804,7 +809,12 @@ def test_cuda_stream_policy_from_env(
     elif runtime == "rapidsmpf":
         with pytest.raises(ValueError, match="The rapidsmpf streaming engine"):
             ConfigOptions.from_polars_engine(engine)
-    else:
+    elif (
+        os.environ.get("CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD", "") != "rapidsmpf"
+        and not rapidsmpf_single_available
+    ):
+        # CUDF_POLARS__EXECUTOR__SHUFFLE_METHOD set in ci/run_cudf_polars_pytests.sh
+        # Would raise since rapidsmpf shuffle requires rapidsmpf to be installed
         config = ConfigOptions.from_polars_engine(engine)
         assert config.cuda_stream_policy == env
 

python/cudf_polars/tests/test_dataframescan.py

@@ -86,7 +86,13 @@ def test_dataframescan_with_decimals():
     POLARS_VERSION_LT_138,
     reason="height parameter added in Polars 1.38",
 )
-def test_dataframescan_zero_width_with_rows():
+def test_dataframescan_zero_width_with_rows(request, using_rapidsmpf):
+    request.applymarker(
+        pytest.mark.xfail(
+            using_rapidsmpf,
+            reason="https://github.com/rapidsai/cudf/issues/21644",
+        )
+    )
     df = pl.LazyFrame(height=5)
     q = df.select(pl.len())
     assert_gpu_result_equal(q)

python/cudf_polars/tests/test_drop_nulls.py

@@ -42,7 +42,15 @@ def test_drop_null(null_data):
     [0, pl.col("a").mean(), pl.col("b")],
     ids=["scalar", "aggregation", "column_expression"],
 )
-def test_fill_null(null_data, value):
+def test_fill_null(null_data, value, using_rapidsmpf, request):
+    request.applymarker(
+        pytest.mark.xfail(
+            condition=using_rapidsmpf
+            and "mean" in str(value)
+            and len(null_data.collect()) == 2,
+            reason="https://github.com/rapidsai/cudf/issues/21721",
+        )
+    )
     q = null_data.select(pl.col("a").fill_null(value))
     assert_gpu_result_equal(q)