diff --git a/ci/run_cudf_polars_experimental_pytests.sh b/ci/run_cudf_polars_experimental_pytests.sh index d58de2b71ac1..3056bd225955 100755 --- a/ci/run_cudf_polars_experimental_pytests.sh +++ b/ci/run_cudf_polars_experimental_pytests.sh @@ -10,18 +10,11 @@ set -euo pipefail # Support invoking outside the script directory cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/ -rapids-logger "Running experimental legacy tests with the 'rapidsmpf' runtime and a 'single' cluster" -timeout 10m python -m pytest --cache-clear "$@" "tests" \ - --executor streaming \ - --cluster single \ - --runtime rapidsmpf \ - --blocksize-mode small +echo "Running the full cudf-polars test suite with both the in-memory and spmd engine" +timeout 10m python -m pytest --cache-clear "$@" tests --ignore=tests/experimental/legacy -rapids-logger "Running experimental legacy tests with the 'rapidsmpf' runtime and a 'distributed' cluster" +echo "Running experimental legacy tests with the 'rapidsmpf' runtime and a 'distributed' cluster" timeout 10m python -m pytest --cache-clear "$@" "tests/experimental/legacy" \ --executor streaming \ --cluster distributed \ --runtime rapidsmpf - -rapids-logger "Running experimental tests" -timeout 10m python -m pytest --cache-clear "$@" tests/experimental --ignore=tests/experimental/legacy diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index 93a267d23085..96e77c4b0389 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -8,9 +8,4 @@ set -euo pipefail # Support invoking run_cudf_polars_pytests.sh outside the script directory cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/ -# Run all non-experimental tests using both the in-memory and streaming executor. -IGNORE_EXPERIMENTAL="--ignore=tests/experimental/" -python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor in-memory -python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor streaming -python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor streaming \ - --blocksize-mode small +python -m pytest --cache-clear "$@" tests --ignore=tests/experimental diff --git a/python/cudf_polars/cudf_polars/testing/asserts.py b/python/cudf_polars/cudf_polars/testing/asserts.py index 413205ec6e51..9f0953cd4dff 100644 --- a/python/cudf_polars/cudf_polars/testing/asserts.py +++ b/python/cudf_polars/cudf_polars/testing/asserts.py @@ -6,18 +6,18 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any import polars as pl from polars import GPUEngine from polars.testing.asserts import assert_frame_equal from cudf_polars.dsl.translate import Translator -from cudf_polars.utils.config import ConfigOptions, StreamingFallbackMode +from cudf_polars.utils.config import ConfigOptions from cudf_polars.utils.versions import POLARS_VERSION_LT_1323 if TYPE_CHECKING: - from cudf_polars.typing import OptimizationArgs + from cudf_polars.typing import CollectKwargs __all__: list[str] = [ @@ -32,16 +32,15 @@ DEFAULT_EXECUTOR = "in-memory" DEFAULT_RUNTIME = "tasks" DEFAULT_CLUSTER = "single" -DEFAULT_BLOCKSIZE_MODE: Literal["small", "default"] = "default" def assert_gpu_result_equal( lazydf: pl.LazyFrame, *, engine: GPUEngine | None = None, - collect_kwargs: dict[OptimizationArgs, bool] | None = None, - polars_collect_kwargs: dict[OptimizationArgs, bool] | None = None, - cudf_collect_kwargs: dict[OptimizationArgs, bool] | None = None, + collect_kwargs: CollectKwargs | None = None, + polars_collect_kwargs: CollectKwargs | None = None, + cudf_collect_kwargs: CollectKwargs | None = None, check_row_order: bool = True, check_column_order: bool = True, check_dtypes: bool = True, @@ -50,7 +49,6 @@ def assert_gpu_result_equal( atol: float = 1e-08, categorical_as_str: bool = False, executor: str | None = None, - blocksize_mode: Literal["small", "default"] | None = None, ) -> None: """ Assert that collection of a lazyframe on GPU produces correct results. @@ -91,12 +89,6 @@ def assert_gpu_result_equal( executor The executor configuration to pass to `GPUEngine`. If not specified uses the module level `Executor` attribute. - blocksize_mode - The "mode" to use for choosing the blocksize for the streaming executor. - If not specified, uses the module level ``DEFAULT_BLOCKSIZE_MODE`` attribute. - Set to "small" to configure small values for ``max_rows_per_partition`` - and ``target_partition_size``, which will typically cause many partitions - to be created while executing the query. Raises ------ @@ -105,7 +97,7 @@ def assert_gpu_result_equal( NotImplementedError If GPU collection failed in some way. """ - engine = engine or get_default_engine(executor, blocksize_mode) + engine = engine or get_default_engine(executor) final_polars_collect_kwargs, final_cudf_collect_kwargs = _process_kwargs( collect_kwargs, polars_collect_kwargs, cudf_collect_kwargs ) @@ -184,7 +176,6 @@ def assert_ir_translation_raises(q: pl.LazyFrame, *exceptions: type[Exception]) def get_default_engine( executor: str | None = None, - blocksize_mode: Literal["small", "default"] | None = None, ) -> GPUEngine: """ Get the default engine used for testing. @@ -194,12 +185,6 @@ def get_default_engine( executor The executor configuration to pass to `GPUEngine`. If not specified uses the module level `Executor` attribute. - blocksize_mode - The "mode" to use for choosing the blocksize for the streaming executor. - If not specified, uses the module level ``DEFAULT_BLOCKSIZE_MODE`` attribute. - Set to "small" to configure small values for ``max_rows_per_partition`` - and ``target_partition_size``, which will typically cause many partitions - to be created while executing the query. Returns ------- @@ -217,14 +202,6 @@ def get_default_engine( executor_options["cluster"] = DEFAULT_CLUSTER executor_options["runtime"] = DEFAULT_RUNTIME - blocksize_mode = blocksize_mode or DEFAULT_BLOCKSIZE_MODE - - if blocksize_mode == "small": # pragma: no cover - executor_options["max_rows_per_partition"] = 4 - executor_options["target_partition_size"] = 10 - # We expect many tests to fall back, so silence the warnings - executor_options["fallback_mode"] = StreamingFallbackMode.SILENT - return GPUEngine( raise_on_fail=True, executor=executor, @@ -233,10 +210,10 @@ def get_default_engine( def _process_kwargs( - collect_kwargs: dict[OptimizationArgs, bool] | None, - polars_collect_kwargs: dict[OptimizationArgs, bool] | None, - cudf_collect_kwargs: dict[OptimizationArgs, bool] | None, -) -> tuple[dict[OptimizationArgs, bool], dict[OptimizationArgs, bool]]: + collect_kwargs: CollectKwargs | None, + polars_collect_kwargs: CollectKwargs | None, + cudf_collect_kwargs: CollectKwargs | None, +) -> tuple[CollectKwargs, CollectKwargs]: if collect_kwargs is None: collect_kwargs = {} final_polars_collect_kwargs = collect_kwargs.copy() @@ -253,9 +230,9 @@ def assert_collect_raises( *, polars_except: type[Exception] | tuple[type[Exception], ...], cudf_except: type[Exception] | tuple[type[Exception], ...], - collect_kwargs: dict[OptimizationArgs, bool] | None = None, - polars_collect_kwargs: dict[OptimizationArgs, bool] | None = None, - cudf_collect_kwargs: dict[OptimizationArgs, bool] | None = None, + collect_kwargs: CollectKwargs | None = None, + polars_collect_kwargs: CollectKwargs | None = None, + cudf_collect_kwargs: CollectKwargs | None = None, ) -> None: """ Assert that collecting the result of a query raises the expected exceptions. @@ -350,7 +327,6 @@ def assert_sink_result_equal( read_kwargs: dict | None = None, write_kwargs: dict | None = None, executor: str | None = None, - blocksize_mode: Literal["small", "default"] | None = None, ) -> None: """ Assert that writing a LazyFrame via sink produces the same output. @@ -371,12 +347,6 @@ def assert_sink_result_equal( executor The executor configuration to pass to `GPUEngine`. If not specified uses the module level `Executor` attribute. - blocksize_mode - The "mode" to use for choosing the blocksize for the streaming executor. - If not specified, uses the module level ``DEFAULT_BLOCKSIZE_MODE`` attribute. - Set to "small" to configure small values for ``max_rows_per_partition`` - and ``target_partition_size``, which will typically cause many partitions - to be created while executing the query. Raises ------ @@ -385,7 +355,7 @@ def assert_sink_result_equal( ValueError If the file extension is not one of the supported formats. """ - engine = engine or get_default_engine(executor, blocksize_mode) + engine = engine or get_default_engine(executor) path = Path(path) read_kwargs = read_kwargs or {} write_kwargs = write_kwargs or {} diff --git a/python/cudf_polars/cudf_polars/typing/__init__.py b/python/cudf_polars/cudf_polars/typing/__init__.py index af59a5119481..ab12787a124b 100644 --- a/python/cudf_polars/cudf_polars/typing/__init__.py +++ b/python/cudf_polars/cudf_polars/typing/__init__.py @@ -33,6 +33,7 @@ __all__: list[str] = [ "ClosedInterval", + "CollectKwargs", "ColumnHeader", "ColumnOptions", "DataFrameHeader", @@ -145,8 +146,11 @@ def set_udf( "comm_subexpr_elim", "cluster_with_columns", "no_optimization", + "optimizations", ] +CollectKwargs: TypeAlias = dict[OptimizationArgs, bool | pl.QueryOptFlags] + U_contra = TypeVar("U_contra", bound=Hashable, contravariant=True) V_co = TypeVar("V_co", covariant=True) diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index de923c60cefd..cf0065f311f4 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -3,10 +3,21 @@ from __future__ import annotations import importlib.util +from typing import TYPE_CHECKING, Any, Literal import pytest +import polars as pl + import cudf_polars.callback +from cudf_polars.utils.config import StreamingFallbackMode + +if TYPE_CHECKING: + from collections.abc import Generator + + from rapidsmpf.communicator.communicator import Communicator + + from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine @pytest.fixture(params=[False, True], ids=["no_nulls", "nulls"], scope="session") @@ -29,8 +40,163 @@ def clear_memory_resource_cache(): @pytest.fixture -def using_rapidsmpf(): - return cudf_polars.testing.asserts.DEFAULT_RUNTIME == "rapidsmpf" +def using_streaming_engine(engine: pl.GPUEngine) -> bool: + """True when the active ``engine`` fixture is a :class:`StreamingEngine`.""" + try: + from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + + return isinstance(engine, StreamingEngine) + except ImportError: + return False + + +@pytest.fixture(autouse=True) +def _skip_unless_spmd(request: pytest.FixtureRequest) -> None: + """Skip tests in SPMD multi-rank mode unless marked with ``pytest.mark.spmd``.""" + # Do not use `pytest.importorskip` here: this fixture is autouse, so an + # import-based skip would skip every test in the suite on environments + # without rapidsmpf (e.g. the coverage CI job), masking real coverage. + # We only want to gate the nranks>1 check on rapidsmpf being available. + if importlib.util.find_spec("rapidsmpf") is None: + return + + from rapidsmpf.bootstrap import get_nranks, is_running_with_rrun + + if ( + is_running_with_rrun() + and get_nranks() > 1 + and not request.node.get_closest_marker("spmd") + ): + pytest.skip("skip: SPMD nranks > 1 (mark with pytest.mark.spmd to run)") + + +@pytest.fixture(scope="session") +def spmd_comm() -> Communicator: + """Session-scoped communicator — bootstrapped once and shared across all tests. + + Sharing a single communicator avoids the file-based bootstrap race that can + cause hangs when ``create_ucxx_comm()`` is called repeatedly in the same + ``rrun`` session (stale barrier files / stale ``ucxx_root_address`` KV entry). + """ + pytest.importorskip("rapidsmpf") + from rapidsmpf import bootstrap + from rapidsmpf.communicator.single import new_communicator as single_communicator + from rapidsmpf.config import Options, get_environment_variables + from rapidsmpf.progress_thread import ProgressThread + + if bootstrap.is_running_with_rrun(): + return bootstrap.create_ucxx_comm( + progress_thread=ProgressThread(), + type=bootstrap.BackendType.AUTO, + ) + return single_communicator(Options(get_environment_variables()), ProgressThread()) + + +@pytest.fixture +def blocksize_mode(request: pytest.FixtureRequest) -> Literal["default", "small"]: + """Blocksize mode for the streaming executor. + + Defaults to ``"default"``. Tests can override this via ``indirect`` + parametrization with ``["default", "small"]`` to run under both the + standard and small-partition configurations. In addition, the + ``engine="spmd-small"`` variant of the ``engine`` fixture implicitly + selects ``"small"`` mode so that every streaming-engine test exercises + tiny-partition / fallback paths without per-test opt-in. Explicit + indirect parametrization always wins over the implicit engine-derived + value. + """ + if hasattr(request, "param"): + return request.param + callspec = getattr(request.node, "callspec", None) + if callspec is not None and callspec.params.get("engine") == "spmd-small": + return "small" + return "default" + + +@pytest.fixture +def streaming_engine( + request: pytest.FixtureRequest, + spmd_comm: Communicator, + blocksize_mode: Literal["default", "small"], +) -> Generator[StreamingEngine, None, None]: + """Yield an :class:`SPMDEngine` configured for streaming-only tests. + + Options can be overridden via ``indirect`` parametrization by passing + a dict with any of the keys ``"executor_options"``, + ``"engine_options"``, or ``"rapidsmpf_options"``. + """ + from rapidsmpf.config import Options + + from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + + params: dict[str, Any] = getattr(request, "param", {}) or {} + executor_options: dict[str, Any] = { + "max_rows_per_partition": 50, + "dynamic_planning": {}, + "target_partition_size": 1_000_000, + } + if blocksize_mode == "small": + executor_options.update( + max_rows_per_partition=4, + target_partition_size=10, + # We expect many tests to fall back, so silence the warnings + fallback_mode=StreamingFallbackMode.SILENT, + ) + executor_options.update(params.get("executor_options", {})) + rapidsmpf_options = ( + Options(params.get("rapidsmpf_options")) + if "rapidsmpf_options" in params + else None + ) + engine_options: dict[str, Any] = {"raise_on_fail": True} + engine_options.update(params.get("engine_options", {})) + with SPMDEngine( + comm=spmd_comm, + rapidsmpf_options=rapidsmpf_options, + executor_options=executor_options, + engine_options=engine_options, + ) as engine: + yield engine + + +_ENGINE_PARAMS = ["in-memory"] +if importlib.util.find_spec("rapidsmpf") is not None: + _ENGINE_PARAMS.extend(["spmd", "spmd-small"]) + + +@pytest.fixture(params=_ENGINE_PARAMS) +def engine( + request: pytest.FixtureRequest, +) -> Generator[pl.GPUEngine, None, None]: + """Yield a :class:`polars.GPUEngine` for each engine variant under test. + + Use this fixture for tests that support any ``GPUEngine``. The test runs + once per available variant: always with the in-memory executor, and, if + ``rapidsmpf`` is installed, a streaming :class:`SPMDEngine` at the + default blocksize (``"spmd"``) and a second streaming engine with + tiny-partition / silent-fallback settings (``"spmd-small"``). The + ``"spmd-small"`` variant forces ``blocksize_mode="small"`` via the + ``blocksize_mode`` fixture, so every test using ``engine`` exercises + multi-partition paths for free. + + For tests that require a ``StreamingEngine``, use the ``streaming_engine`` + fixture instead. + """ + if request.param == "in-memory": + yield pl.GPUEngine(executor="in-memory", raise_on_fail=True) + else: + yield request.getfixturevalue("streaming_engine") + + +@pytest.fixture +def engine_raise_on_fail() -> pl.GPUEngine: + """Yield a default :class:`polars.GPUEngine` with ``raise_on_fail=True``. + + Intended for error-path tests that assert specific exceptions propagate + from ``.collect()``. Uses the default (in-memory) executor so errors are + not wrapped by a streaming task group. + """ + return pl.GPUEngine(raise_on_fail=True) def pytest_addoption(parser): @@ -58,21 +224,28 @@ def pytest_addoption(parser): help="Cluster to use for 'streaming' executor.", ) - parser.addoption( - "--blocksize-mode", - action="store", - default="default", - choices=("small", "default"), - help=( - "Blocksize to use for 'streaming' executor. Set to 'small' " - "to run most tests with multiple partitions." - ), - ) - def pytest_configure(config): import cudf_polars.testing.asserts + config.addinivalue_line( + "markers", + "skip_on_streaming_engine(reason): skip the test when the `engine` " + "fixture resolves to a streaming engine variant (e.g. 'spmd'). " + "Use for tests exercising operations that have no multi-partition " + "implementation.", + ) + + # Ray's internal subprocess management leaks `/dev/null` file handles, and + # distributed's shutdown leaves unclosed sockets. Under Python 3.14 + + # pytest 9, these surface as unraisable `ResourceWarning`s and — combined + # with `filterwarnings = ["error", ...]` in pyproject.toml — fail + # otherwise-unrelated tests when the GC finalizer happens to fire during + # them. With `pytest-xdist --dist=worksteal`, the leak can land in any + # test that shares a worker with a ray/dask test, so the suppression must + # apply globally rather than per-module. + config.addinivalue_line("filterwarnings", "ignore::ResourceWarning") + if ( config.getoption("--cluster") == "distributed" and config.getoption("--executor") != "streaming" @@ -91,6 +264,23 @@ def pytest_configure(config): cudf_polars.testing.asserts.DEFAULT_EXECUTOR = config.getoption("--executor") cudf_polars.testing.asserts.DEFAULT_RUNTIME = config.getoption("--runtime") cudf_polars.testing.asserts.DEFAULT_CLUSTER = config.getoption("--cluster") - cudf_polars.testing.asserts.DEFAULT_BLOCKSIZE_MODE = config.getoption( - "--blocksize-mode" - ) + + +def pytest_collection_modifyitems(items): + """Apply ``skip_on_streaming_engine`` markers to parametrized ``engine`` items.""" + for item in items: + marker = item.get_closest_marker("skip_on_streaming_engine") + if marker is None: + continue + callspec = getattr(item, "callspec", None) + if callspec is None: + continue + engine_param = callspec.params.get("engine") + if engine_param is None or engine_param == "in-memory": + continue + reason = ( + marker.args[0] + if marker.args + else marker.kwargs.get("reason", "unsupported on streaming engine") + ) + item.add_marker(pytest.mark.skip(reason=reason)) diff --git a/python/cudf_polars/tests/containers/test_dataframe.py b/python/cudf_polars/tests/containers/test_dataframe.py index 67f393c9b446..5006fcaa9b4a 100644 --- a/python/cudf_polars/tests/containers/test_dataframe.py +++ b/python/cudf_polars/tests/containers/test_dataframe.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 @@ -176,14 +176,14 @@ def test_sorted_flags_preserved(with_nulls, nulls_last): assert df.flags == gf.to_polars().flags -def test_empty_name_roundtrips_overlap(): +def test_empty_name_roundtrips_overlap(engine: pl.GPUEngine): df = pl.LazyFrame({"": [1, 2, 3], "column_0": [4, 5, 6]}) - assert_gpu_result_equal(df) + assert_gpu_result_equal(df, engine=engine) -def test_empty_name_roundtrips_no_overlap(): +def test_empty_name_roundtrips_no_overlap(engine: pl.GPUEngine): df = pl.LazyFrame({"": [1, 2, 3], "b": [4, 5, 6]}) - assert_gpu_result_equal(df) + assert_gpu_result_equal(df, engine=engine) @pytest.mark.parametrize( diff --git a/python/cudf_polars/tests/experimental/conftest.py b/python/cudf_polars/tests/experimental/conftest.py deleted file mode 100644 index ad05928ab409..000000000000 --- a/python/cudf_polars/tests/experimental/conftest.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""Configuration for StreamingEngine tests.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -import pytest -from rapidsmpf import bootstrap -from rapidsmpf.bootstrap import get_nranks, is_running_with_rrun -from rapidsmpf.communicator.single import new_communicator as single_communicator -from rapidsmpf.config import Options, get_environment_variables -from rapidsmpf.progress_thread import ProgressThread - -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine - -if TYPE_CHECKING: - from collections.abc import Generator - - from rapidsmpf.communicator.communicator import Communicator - - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine - - -@pytest.fixture(autouse=True) -def _skip_unless_spmd(request: pytest.FixtureRequest) -> None: - """Skip tests in SPMD multi-rank mode unless marked with ``pytest.mark.spmd``.""" - if ( - is_running_with_rrun() - and get_nranks() > 1 - and not request.node.get_closest_marker("spmd") - ): - pytest.skip("skip: SPMD nranks > 1 (mark with pytest.mark.spmd to run)") - - -@pytest.fixture(scope="session") -def spmd_comm() -> Communicator: - """Session-scoped communicator — bootstrapped once and shared across all tests. - - Sharing a single communicator avoids the file-based bootstrap race that can - cause hangs when ``create_ucxx_comm()`` is called repeatedly in the same - ``rrun`` session (stale barrier files / stale ``ucxx_root_address`` KV entry). - """ - if bootstrap.is_running_with_rrun(): - return bootstrap.create_ucxx_comm( - progress_thread=ProgressThread(), - type=bootstrap.BackendType.AUTO, - ) - return single_communicator(Options(get_environment_variables()), ProgressThread()) - - -@pytest.fixture -def engine( - request: pytest.FixtureRequest, - spmd_comm: Communicator, -) -> Generator[StreamingEngine, None, None]: - """Yield a :class:`~cudf_polars.experimental.rapidsmpf.frontend.spmd.SPMDEngine`. - - Default executor options enable dynamic planning with sensible partition - sizes. Override options via ``indirect`` parametrization by passing a dict - with ``"executor_options"`` and/or ``"engine_options"`` keys: - - .. code-block:: python - - @pytest.mark.parametrize( - "engine", - [{"executor_options": {"target_partition_size": 100_000_000}}], - indirect=True, - ) - def test_foo(engine): ... - """ - params: dict[str, Any] = getattr(request, "param", {}) - executor_options = { - "max_rows_per_partition": 50, - "dynamic_planning": {}, - "target_partition_size": 1_000_000, - **params.get("executor_options", {}), - } - - rapidsmpf_options = ( - Options(params.get("rapidsmpf_options")) - if "rapidsmpf_options" in params - else None - ) - - with SPMDEngine( - comm=spmd_comm, - rapidsmpf_options=rapidsmpf_options, - executor_options=executor_options, - engine_options=params.get("engine_options", {}), - ) as engine: - yield engine diff --git a/python/cudf_polars/tests/experimental/test_agg.py b/python/cudf_polars/tests/experimental/test_agg.py index 81cd6757354d..beeeab51ed82 100644 --- a/python/cudf_polars/tests/experimental/test_agg.py +++ b/python/cudf_polars/tests/experimental/test_agg.py @@ -24,7 +24,7 @@ def decimal_df() -> pl.LazyFrame: ) -def test_decimal_aggs(decimal_df: pl.LazyFrame, engine) -> None: +def test_decimal_aggs(decimal_df: pl.LazyFrame, streaming_engine) -> None: q = decimal_df.with_columns( sum=pl.col("a").sum(), min=pl.col("a").min(), @@ -32,10 +32,10 @@ def test_decimal_aggs(decimal_df: pl.LazyFrame, engine) -> None: mean=pl.col("a").mean(), median=pl.col("a").median(), ) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) -def test_mean_all_null(engine): +def test_mean_all_null(streaming_engine): lf = pl.LazyFrame({"a": [None, None]}, schema={"a": pl.Float64}) q = lf.select(pl.col("a").mean()) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) diff --git a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py b/python/cudf_polars/tests/experimental/test_all_gather_host_data.py index b308c8e8cafc..df71039b131d 100644 --- a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py +++ b/python/cudf_polars/tests/experimental/test_all_gather_host_data.py @@ -34,20 +34,20 @@ def _struct(rank: int) -> bytes: @pytest.mark.parametrize("make_data", [_empty, _text, _bytearray, _struct]) -def test_all_gather_host_data(engine, make_data) -> None: +def test_all_gather_host_data(streaming_engine, make_data) -> None: """Each rank sends rank-specific data; results are correct and ordered.""" - comm = engine.comm - br = engine.context.br() + comm = streaming_engine.comm + br = streaming_engine.context.br() result = all_gather_host_data(comm, br, op_id=0, data=make_data(comm.rank)) assert len(result) == comm.nranks for i, item in enumerate(result): assert item == bytes(make_data(i)) -def test_gather_cluster_info(engine) -> None: +def test_gather_cluster_info(streaming_engine) -> None: """SPMDEngine.gather_cluster_info returns ClusterInfo for each rank.""" - infos = engine.gather_cluster_info() - assert len(infos) == engine.nranks + infos = streaming_engine.gather_cluster_info() + assert len(infos) == streaming_engine.nranks for info in infos: assert isinstance(info, ClusterInfo) assert info.pid > 0 @@ -57,9 +57,9 @@ def test_gather_cluster_info(engine) -> None: ) assert isinstance(info.gpu_uuid, str) # Each rank runs in its own process. - assert len({info.pid for info in infos}) == engine.nranks + assert len({info.pid for info in infos}) == streaming_engine.nranks # Without allow_gpu_sharing, all UUIDs must be unique (enforced at init). - assert len({info.gpu_uuid for info in infos}) == engine.nranks + assert len({info.gpu_uuid for info in infos}) == streaming_engine.nranks def test_cluster_info_cuda_visible_devices(monkeypatch) -> None: @@ -75,10 +75,10 @@ def test_cluster_info_cuda_visible_devices_unset(monkeypatch) -> None: @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"engine_options": {"allow_gpu_sharing": True}}], indirect=True, ) -def test_allow_gpu_sharing(engine) -> None: +def test_allow_gpu_sharing(streaming_engine) -> None: """Engine init succeeds with allow_gpu_sharing=True.""" - assert engine.nranks >= 1 + assert streaming_engine.nranks >= 1 diff --git a/python/cudf_polars/tests/experimental/test_allgather.py b/python/cudf_polars/tests/experimental/test_allgather.py index 464db3feb98a..e2c241b09075 100644 --- a/python/cudf_polars/tests/experimental/test_allgather.py +++ b/python/cudf_polars/tests/experimental/test_allgather.py @@ -52,8 +52,8 @@ async def _test_allgather(engine) -> None: assert col.type().id().value == plc.types.TypeId.INT32.value -def test_allgather(engine) -> None: - asyncio.run(_test_allgather(engine)) +def test_allgather(streaming_engine) -> None: + asyncio.run(_test_allgather(streaming_engine)) async def _test_allgather_reduce(engine) -> None: @@ -70,5 +70,5 @@ async def _test_allgather_reduce(engine) -> None: assert results == (10, 20, 30) # Single rank, so sums are just the local values -def test_allgather_reduce(engine) -> None: - asyncio.run(_test_allgather_reduce(engine)) +def test_allgather_reduce(streaming_engine) -> None: + asyncio.run(_test_allgather_reduce(streaming_engine)) diff --git a/python/cudf_polars/tests/experimental/test_dask.py b/python/cudf_polars/tests/experimental/test_dask.py index 0d4631d54657..d923edd37cf9 100644 --- a/python/cudf_polars/tests/experimental/test_dask.py +++ b/python/cudf_polars/tests/experimental/test_dask.py @@ -36,8 +36,6 @@ def engine() -> Iterator[DaskEngine]: pytestmark = [ - # distributed's shutdown leaves unclosed sockets; suppress the noise. - pytest.mark.filterwarnings("ignore::ResourceWarning"), pytest.mark.skipif( is_running_with_rrun(), reason="DaskEngine must not be created from within an rrun cluster", diff --git a/python/cudf_polars/tests/experimental/test_dataframescan.py b/python/cudf_polars/tests/experimental/test_dataframescan.py index 422c1336b064..1ef4abacb7b5 100644 --- a/python/cudf_polars/tests/experimental/test_dataframescan.py +++ b/python/cudf_polars/tests/experimental/test_dataframescan.py @@ -33,16 +33,16 @@ def df(): @pytest.mark.parametrize( - "max_rows_per_partition,engine", + "max_rows_per_partition,streaming_engine", [ (1_000, {"executor_options": {"max_rows_per_partition": 1_000}}), (1_000_000, {"executor_options": {"max_rows_per_partition": 1_000_000}}), ], - indirect=["engine"], + indirect=["streaming_engine"], ) -def test_parallel_dataframescan(df, max_rows_per_partition, engine): +def test_parallel_dataframescan(df, max_rows_per_partition, streaming_engine): total_row_count = len(df.collect()) - assert_gpu_result_equal(df, engine=engine) + assert_gpu_result_equal(df, engine=streaming_engine) # Check partitioning (throwaway engine — no cluster/runtime needed) _engine = pl.GPUEngine( @@ -60,13 +60,13 @@ def test_parallel_dataframescan(df, max_rows_per_partition, engine): @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"max_rows_per_partition": 1_000}}], indirect=True, ) -def test_dataframescan_concat(df, engine): +def test_dataframescan_concat(df, streaming_engine): df2 = pl.concat([df, df]) - assert_gpu_result_equal(df2, engine=engine) + assert_gpu_result_equal(df2, engine=streaming_engine) def test_join_in_memory_lazy_stable_id_pickle(): diff --git a/python/cudf_polars/tests/experimental/test_distinct.py b/python/cudf_polars/tests/experimental/test_distinct.py index 4982f5b9069a..d2a7d42527e1 100644 --- a/python/cudf_polars/tests/experimental/test_distinct.py +++ b/python/cudf_polars/tests/experimental/test_distinct.py @@ -25,7 +25,7 @@ def df() -> pl.LazyFrame: @pytest.mark.parametrize("subset", [None, ("y",), ("x", "y")]) @pytest.mark.parametrize("keep", ["any", "none"]) -def test_dynamic_distinct_basic(df, engine, subset, keep): +def test_dynamic_distinct_basic(df, streaming_engine, subset, keep): """Test dynamic distinct with various subset and keep options.""" q = df.unique(subset=subset, keep=keep, maintain_order=False) @@ -34,50 +34,50 @@ def test_dynamic_distinct_basic(df, engine, subset, keep): if keep == "any" and subset: q = q.select(*(pl.col(col) for col in subset)) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 100_000_000}}], indirect=True, ) -def test_dynamic_distinct_tree_strategy(df, engine): +def test_dynamic_distinct_tree_strategy(df, streaming_engine): """Test that small output uses tree reduction (high target_partition_size).""" subset = ("y",) q = df.unique(subset=subset, keep="any", maintain_order=False) # With keep="any", non-subset columns are non-deterministic, so only check subset q = q.select(*(pl.col(col) for col in subset)) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 1000}}], indirect=True, ) -def test_dynamic_distinct_shuffle_strategy(engine): +def test_dynamic_distinct_shuffle_strategy(streaming_engine): """Test that large output uses shuffle (low target_partition_size).""" df = pl.LazyFrame({"x": range(1000), "y": range(1000)}) q = df.unique(subset=None, keep="any", maintain_order=False) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_distinct_all_duplicates(engine): +def test_dynamic_distinct_all_duplicates(streaming_engine): """Test dynamic distinct where all rows are duplicates.""" df = pl.LazyFrame({"x": [1, 1, 1, 1], "y": [2, 2, 2, 2]}) q = df.unique() - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_distinct_single_row(engine): +def test_dynamic_distinct_single_row(streaming_engine): """Test dynamic distinct on single-row DataFrame.""" df = pl.LazyFrame({"x": [1], "y": [2]}) q = df.unique() - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_distinct_chunkwise_after_groupby(engine): +def test_dynamic_distinct_chunkwise_after_groupby(streaming_engine): """Test distinct after group_by is handled chunkwise.""" df = pl.LazyFrame( { @@ -87,4 +87,4 @@ def test_dynamic_distinct_chunkwise_after_groupby(engine): ) # Groupby partitions data by "key", then unique on "key" should be chunkwise q = df.group_by("key").agg(pl.col("value").sum()).unique(subset=("key",)) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) diff --git a/python/cudf_polars/tests/experimental/test_groupby.py b/python/cudf_polars/tests/experimental/test_groupby.py index 11652f0c312d..16739f68d743 100644 --- a/python/cudf_polars/tests/experimental/test_groupby.py +++ b/python/cudf_polars/tests/experimental/test_groupby.py @@ -34,64 +34,64 @@ def df() -> pl.LazyFrame: @pytest.mark.parametrize("keys", [("key",), ("key", "key2")]) @pytest.mark.parametrize("agg", ["sum", "mean", "len", "min", "max"]) -def test_dynamic_groupby_basic(df, engine, keys, agg): +def test_dynamic_groupby_basic(df, streaming_engine, keys, agg): """Test dynamic groupby with various key and agg combinations.""" expr = getattr(pl.col("value"), agg)() q = df.group_by(*keys).agg(expr) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 100_000_000}}], indirect=True, ) -def test_dynamic_groupby_tree_strategy(df, engine): +def test_dynamic_groupby_tree_strategy(df, streaming_engine): """Test that small output uses tree reduction (high target_partition_size).""" q = df.group_by("key2").agg(pl.col("value").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 1000}}], indirect=True, ) -def test_dynamic_groupby_shuffle_strategy(engine): +def test_dynamic_groupby_shuffle_strategy(streaming_engine): """Test that large output uses shuffle (low target_partition_size).""" df = pl.LazyFrame({"key": range(1000), "value": range(1000)}) q = df.group_by("key").agg(pl.col("value").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_groupby_single_group(engine): +def test_dynamic_groupby_single_group(streaming_engine): """Test dynamic groupby where all rows have the same key.""" df = pl.LazyFrame({"key": [1] * 100, "value": range(100)}) q = df.group_by("key").agg(pl.col("value").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_groupby_multiple_aggs(df, engine): +def test_dynamic_groupby_multiple_aggs(df, streaming_engine): """Test dynamic groupby with multiple aggregations.""" q = df.group_by("key").agg( pl.col("value").sum().alias("value_sum"), pl.col("value").mean().alias("value_mean"), pl.col("value2").min().alias("value2_min"), ) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_groupby_maintain_order(df, engine): +def test_dynamic_groupby_maintain_order(df, streaming_engine): """Test dynamic groupby with maintain_order=True.""" q = df.group_by("key", maintain_order=True).agg(pl.col("value").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_dynamic_groupby_single_row(engine): +def test_dynamic_groupby_single_row(streaming_engine): """Test dynamic groupby on single-row DataFrame.""" df = pl.LazyFrame({"key": [1], "value": [42]}) q = df.group_by("key").agg(pl.col("value").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) # --------------------------------------------------------------------------- @@ -101,31 +101,31 @@ def test_dynamic_groupby_single_row(engine): @pytest.mark.parametrize("op", ["sum", "mean", "len"]) @pytest.mark.parametrize("keys", [("y",), ("y", "z")]) -def test_groupby(df, engine, op, keys): +def test_groupby(df, streaming_engine, op, keys): q = getattr(df.group_by(*keys), op)() - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize("op", ["sum", "mean", "len"]) @pytest.mark.parametrize("keys", [("y",), ("y", "z")]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"max_rows_per_partition": int(1e9)}}], indirect=True, ) -def test_groupby_single_partitions(df, engine, op, keys): +def test_groupby_single_partitions(df, streaming_engine, op, keys): q = getattr(df.group_by(*keys), op)() - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( "op", ["sum", "mean", "len", "count", "min", "max", "n_unique", "std", "var"] ) @pytest.mark.parametrize("keys", [("y",), ("y", "z")]) -def test_groupby_agg(df, engine, op, keys): +def test_groupby_agg(df, streaming_engine, op, keys): agg = getattr(pl.col("x"), op)() q = df.group_by(*keys).agg(agg) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize("ddof", [0, 2, 50]) @@ -136,16 +136,16 @@ def test_groupby_std_var_ddof(df, engine, agg, ddof): @pytest.mark.parametrize( - "fallback_mode,engine", + "fallback_mode,streaming_engine", [ ("silent", {"executor_options": {"fallback_mode": "silent"}}), ("raise", {"executor_options": {"fallback_mode": "raise"}}), ("warn", {"executor_options": {"fallback_mode": "warn"}}), ("foo", {"executor_options": {"fallback_mode": "foo"}}), ], - indirect=["engine"], + indirect=["streaming_engine"], ) -def test_groupby_fallback(df, fallback_mode, engine): +def test_groupby_fallback(df, fallback_mode, streaming_engine): match = "Failed to decompose groupby aggs" q = df.group_by("y").median() @@ -165,12 +165,12 @@ def test_groupby_fallback(df, fallback_mode, engine): else: ctx = pytest.warns(UserWarning, match=match) with ctx: - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_groupby_agg_literal(df, engine): +def test_groupby_agg_literal(df, streaming_engine): q = df.group_by("y").agg(1) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( @@ -182,9 +182,9 @@ def test_groupby_agg_literal(df, engine): pl.max("x") + 1, ], ) -def test_groupby_agg_binop(df: pl.LazyFrame, engine, op: pl.Expr) -> None: +def test_groupby_agg_binop(df: pl.LazyFrame, streaming_engine, op: pl.Expr) -> None: q = df.group_by("y").agg(op) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( @@ -194,7 +194,7 @@ def test_groupby_agg_binop(df: pl.LazyFrame, engine, op: pl.Expr) -> None: (pl.mean("x"), "x__mean_sum"), ], ) -def test_groupby_agg_duplicate(engine, op: pl.Expr, column_name: str) -> None: +def test_groupby_agg_duplicate(streaming_engine, op: pl.Expr, column_name: str) -> None: # Ensure that the column names we create internally don't collide with # the user's column names. df = pl.LazyFrame( @@ -205,17 +205,17 @@ def test_groupby_agg_duplicate(engine, op: pl.Expr, column_name: str) -> None: } ) q = df.group_by("y").agg(op, pl.min(column_name)) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_groupby_agg_empty(df: pl.LazyFrame, engine) -> None: +def test_groupby_agg_empty(df: pl.LazyFrame, streaming_engine) -> None: q = df.group_by("y").agg() - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.filterwarnings("ignore:This slice not supported for multiple partitions.") @pytest.mark.parametrize("zlice", [(0, 2), (2, 2), (-2, None)]) -def test_groupby_then_slice(engine, zlice: tuple[int, int]) -> None: +def test_groupby_then_slice(streaming_engine, zlice: tuple[int, int]) -> None: df = pl.LazyFrame( { "x": [0, 1, 2, 3] * 2, @@ -223,10 +223,10 @@ def test_groupby_then_slice(engine, zlice: tuple[int, int]) -> None: } ) q = df.group_by("y", maintain_order=True).max().slice(*zlice) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) -def test_groupby_on_equality(engine) -> None: +def test_groupby_on_equality(streaming_engine) -> None: # See: https://github.com/rapidsai/cudf/issues/19152 df = pl.LazyFrame( { @@ -236,7 +236,7 @@ def test_groupby_on_equality(engine) -> None: } ) q = df.group_by(pl.col("key1") == pl.col("key2")).agg(pl.col("int32").sum()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( @@ -247,11 +247,11 @@ def test_groupby_on_equality(engine) -> None: ], ) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"max_rows_per_partition": 2}}], indirect=True, ) -def test_mean_partitioned(values: list[int | None], engine) -> None: +def test_mean_partitioned(values: list[int | None], streaming_engine) -> None: df = pl.LazyFrame( { "key1": [1, 1, 2, 2], @@ -259,10 +259,10 @@ def test_mean_partitioned(values: list[int | None], engine) -> None: } ) q = df.group_by("key1").agg(pl.col("uint16_with_null").mean()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) -def test_groupby_literal_key(df, engine): +def test_groupby_literal_key(df, streaming_engine): q = ( df.group_by( pl.lit(True).alias("key"), # noqa: FBT003 @@ -271,7 +271,7 @@ def test_groupby_literal_key(df, engine): .agg(pl.col("x").sum()) .drop("key") ) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) # --------------------------------------------------------------------------- @@ -282,7 +282,7 @@ def test_groupby_literal_key(df, engine): @pytest.mark.parametrize("op", ["sum", "mean", "len", "count"]) @pytest.mark.parametrize("keys", [("y",), ("y", "z")]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -294,30 +294,30 @@ def test_groupby_literal_key(df, engine): ], indirect=True, ) -def test_groupby_agg_config_options(df, op, keys, engine): +def test_groupby_agg_config_options(df, op, keys, streaming_engine): agg = getattr(pl.col("x"), op)() if op in ("sum", "mean"): agg = agg.round(2) # Unary test coverage q = df.group_by(*keys).agg(agg) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 1}}], indirect=True, ) -def test_groupby_count_type_mismatch(df, engine): +def test_groupby_count_type_mismatch(df, streaming_engine): q = df.group_by("key", maintain_order=True).agg(pl.col("value").count()) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 10, "max_rows_per_partition": 5}}], indirect=True, ) -def test_shuffle_reduce_insert_finished_called_on_oom(engine): +def test_shuffle_reduce_insert_finished_called_on_oom(streaming_engine): # Tests that an exception raised inside insert_hash() must not leave the # C++ ShufflerAsync without insert_finished() being called. @@ -329,5 +329,5 @@ def foo(*args, **kwargs): patch.object(ShuffleManager.Inserter, "insert_hash", foo), pytest.raises(ExceptionGroup) as exc_info, ): - df.group_by("a").agg(pl.col("b").sum()).collect(engine=engine) + df.group_by("a").agg(pl.col("b").sum()).collect(engine=streaming_engine) assert any("OOM in insert_hash" in str(e) for e in exc_info.value.exceptions) diff --git a/python/cudf_polars/tests/experimental/test_io_multirank.py b/python/cudf_polars/tests/experimental/test_io_multirank.py index 0c3ad87cbfe5..e12656023046 100644 --- a/python/cudf_polars/tests/experimental/test_io_multirank.py +++ b/python/cudf_polars/tests/experimental/test_io_multirank.py @@ -27,7 +27,6 @@ # variants skip themselves in that environment. pytestmark = [ pytest.mark.spmd, - pytest.mark.filterwarnings("ignore::ResourceWarning"), ] diff --git a/python/cudf_polars/tests/experimental/test_join.py b/python/cudf_polars/tests/experimental/test_join.py index 009a8183e37b..ba7ab665f213 100644 --- a/python/cudf_polars/tests/experimental/test_join.py +++ b/python/cudf_polars/tests/experimental/test_join.py @@ -44,30 +44,30 @@ def right(): @pytest.mark.parametrize("how", ["inner", "left", "right", "full"]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ {"executor_options": {"max_rows_per_partition": 3, "broadcast_join_limit": 2}}, {"executor_options": {"max_rows_per_partition": 5, "broadcast_join_limit": 2}}, ], indirect=True, ) -def test_dynamic_join_how(left, right, engine, how): +def test_dynamic_join_how(left, right, streaming_engine, how): """Dynamic join path: all join types including Right and Full.""" q = left.join(right, on="y", how=how) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize("how", ["right", "full"]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"max_rows_per_partition": 3, "broadcast_join_limit": 2}}], indirect=True, ) -def test_dynamic_join_right_full_reverse(left, right, engine, how): +def test_dynamic_join_right_full_reverse(left, right, streaming_engine, how): """Dynamic join path: Right/Full with reversed left/right (stress ordering).""" # Reverse so "right" frame is larger; exercises right-side preservation q = right.join(left, on="y", how=how) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) # --------------------------------------------------------------------------- @@ -76,23 +76,23 @@ def test_dynamic_join_right_full_reverse(left, right, engine, how): @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"max_rows_per_partition": 2, "broadcast_join_limit": 1}}], indirect=True, ) -def test_join_then_shuffle(left, right, engine): +def test_join_then_shuffle(left, right, streaming_engine): q = left.join(right, on="y", how="inner").select( pl.col("x").sum(), pl.col("xx").mean(), pl.col("y").n_unique(), (pl.col("y") * pl.col("y")).n_unique().alias("y2"), ) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize("reverse", [True, False]) @pytest.mark.parametrize( - "max_rows_per_partition,engine", + "max_rows_per_partition,streaming_engine", [ ( 3, @@ -115,9 +115,9 @@ def test_join_then_shuffle(left, right, engine): }, ), ], - indirect=["engine"], + indirect=["streaming_engine"], ) -def test_join_conditional(reverse, max_rows_per_partition, engine): +def test_join_conditional(reverse, max_rows_per_partition, streaming_engine): left = pl.LazyFrame({"x": range(15), "y": [1, 2, 3] * 5}) right = pl.LazyFrame({"xx": range(9), "yy": [2, 4, 3] * 3}) if reverse: @@ -127,9 +127,9 @@ def test_join_conditional(reverse, max_rows_per_partition, engine): with pytest.warns( UserWarning, match="ConditionalJoin not supported for multiple partitions." ): - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) else: - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) # --------------------------------------------------------------------------- @@ -140,7 +140,7 @@ def test_join_conditional(reverse, max_rows_per_partition, engine): @pytest.mark.parametrize("how", ["inner", "left", "right", "full", "semi", "anti"]) @pytest.mark.parametrize("reverse", [True, False]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ {"executor_options": {"max_rows_per_partition": 1, "broadcast_join_limit": 1}}, {"executor_options": {"max_rows_per_partition": 1, "broadcast_join_limit": 16}}, @@ -168,13 +168,13 @@ def test_join_conditional(reverse, max_rows_per_partition, engine): ], indirect=True, ) -def test_join(left, right, how, reverse, engine): +def test_join(left, right, how, reverse, streaming_engine): if reverse: left, right = right, left q = left.join(right, on="y", how=how) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) # Join again on the same key. # (covers code path that avoids redundant shuffles) @@ -187,12 +187,12 @@ def test_join(left, right, how, reverse, engine): } ) q2 = q.join(right2, left_on="y", right_on="yyy", how=how) - assert_gpu_result_equal(q2, engine=engine, check_row_order=False) + assert_gpu_result_equal(q2, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize("zlice", [(0, 2), (2, 2), (-2, None)]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -204,7 +204,7 @@ def test_join(left, right, how, reverse, engine): ], indirect=True, ) -def test_join_and_slice(zlice, engine): +def test_join_and_slice(zlice, streaming_engine): left = pl.LazyFrame( { "a": [1, 2, 3, 1, None], @@ -226,9 +226,9 @@ def test_join_and_slice(zlice, engine): with pytest.warns( UserWarning, match="This slice not supported for multiple partitions." ): - assert q.collect(engine=engine).height == q.collect().height + assert q.collect(engine=streaming_engine).height == q.collect().height else: - assert q.collect(engine=engine).height == q.collect().height + assert q.collect(engine=streaming_engine).height == q.collect().height # Need sort to match order after a join q = left.join(right, on="a", how="inner").sort(pl.col("a")).slice(*zlice) @@ -237,14 +237,14 @@ def test_join_and_slice(zlice, engine): UserWarning, match="does not support a multi-partition slice with an offset.", ): - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) else: - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) @pytest.mark.parametrize("how", ["inner", "semi", "left", "right"]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -256,19 +256,19 @@ def test_join_and_slice(zlice, engine): ], indirect=True, ) -def test_bloom_filter_join(how, engine): +def test_bloom_filter_join(how, streaming_engine): dim = pl.LazyFrame({"key": range(10), "val": range(10)}) fact = pl.LazyFrame({"key": range(200), "data": range(200)}) left, right = (dim, fact) if how == "right" else (fact, dim) q = left.join(right, on="key", how=how) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( "maintain_order", ["left_right", "right_left", "left", "right"] ) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -280,14 +280,16 @@ def test_bloom_filter_join(how, engine): ], indirect=True, ) -def test_join_maintain_order_fallback_streaming(left, right, maintain_order, engine): +def test_join_maintain_order_fallback_streaming( + left, right, maintain_order, streaming_engine +): q = left.join(right, on="y", how="inner", maintain_order=maintain_order) with pytest.warns( UserWarning, match=r"Join\(maintain_order=.*\) not supported for multiple partitions\.", ): - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/experimental/test_metadata.py b/python/cudf_polars/tests/experimental/test_metadata.py index 8738e05b61a0..898fe917a90b 100644 --- a/python/cudf_polars/tests/experimental/test_metadata.py +++ b/python/cudf_polars/tests/experimental/test_metadata.py @@ -48,7 +48,7 @@ def right() -> pl.LazyFrame: @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -70,16 +70,16 @@ def right() -> pl.LazyFrame: def test_rapidsmpf_join_metadata( left: pl.LazyFrame, right: pl.LazyFrame, - engine, + streaming_engine, ) -> None: - config_options = ConfigOptions.from_polars_engine(engine) + config_options = ConfigOptions.from_polars_engine(streaming_engine) broadcast_join_limit = config_options.executor.broadcast_join_limit q = left.join( right, on="y", how="left", ).filter(pl.col("x") > pl.col("zz")) - ir = Translator(q._ldf.visit(), engine).translate_ir() + ir = Translator(q._ldf.visit(), streaming_engine).translate_ir() left_count = left.collect().height right_count = right.collect().height @@ -349,13 +349,18 @@ def _make_select_ir(engine: pl.GPUEngine, output_columns: tuple[str, ...]): return Select(out_schema, exprs, should_broadcast=False, df=child) -def test_remap_partitioning_select_none_input(engine) -> None: - assert maybe_remap_partitioning(_make_select_ir(engine, ("a", "b")), None) is None +def test_remap_partitioning_select_none_input(streaming_engine) -> None: + assert ( + maybe_remap_partitioning(_make_select_ir(streaming_engine, ("a", "b")), None) + is None + ) -def test_remap_partitioning_select_preserves_keys(engine) -> None: +def test_remap_partitioning_select_preserves_keys(streaming_engine) -> None: part = Partitioning(inter_rank=HashScheme((0, 1), 8), local="inherit") - result = maybe_remap_partitioning(_make_select_ir(engine, ("a", "b")), part) + result = maybe_remap_partitioning( + _make_select_ir(streaming_engine, ("a", "b")), part + ) assert result is not None assert result.inter_rank is not None assert result.inter_rank.column_indices == (0, 1) @@ -363,14 +368,14 @@ def test_remap_partitioning_select_preserves_keys(engine) -> None: assert result.local == "inherit" -def test_remap_partitioning_groupby(engine) -> None: +def test_remap_partitioning_groupby(streaming_engine) -> None: """Hash indices refer to the groupby input child; remap to groupby output columns.""" q = ( pl.LazyFrame({"a": [1], "b": [2], "c": [3]}) .group_by("a", "b") .agg(pl.col("c").sum()) ) - ir = Translator(q._ldf.visit(), engine).translate_ir() + ir = Translator(q._ldf.visit(), streaming_engine).translate_ir() while isinstance(ir, (Select, Projection)): ir = ir.children[0] assert isinstance(ir, GroupBy) @@ -391,9 +396,9 @@ def test_remap_partitioning_groupby(engine) -> None: assert result.local == "inherit" -def test_remap_partitioning_hstack_appends_preserves_keys(engine) -> None: +def test_remap_partitioning_hstack_appends_preserves_keys(streaming_engine) -> None: q = pl.LazyFrame({"a": [1], "b": [2], "c": [3]}) - child = Translator(q._ldf.visit(), engine).translate_ir() + child = Translator(q._ldf.visit(), streaming_engine).translate_ir() d_dtype = DataType(pl.Int64()) hstack = HStack( {**child.schema, "d": d_dtype}, @@ -410,17 +415,17 @@ def test_remap_partitioning_hstack_appends_preserves_keys(engine) -> None: assert result.local == "inherit" -def test_remap_partitioning_select_drops_key(engine) -> None: +def test_remap_partitioning_select_drops_key(streaming_engine) -> None: part = Partitioning(inter_rank=HashScheme((0, 1), 8), local="inherit") - result = maybe_remap_partitioning(_make_select_ir(engine, ("a",)), part) + result = maybe_remap_partitioning(_make_select_ir(streaming_engine, ("a",)), part) assert result is not None assert result.inter_rank is None assert result.local == "inherit" -def test_remap_partitioning_select_renamed_key(engine) -> None: +def test_remap_partitioning_select_renamed_key(streaming_engine) -> None: q = pl.LazyFrame({"a": [1], "b": [2], "c": [3]}) - child = Translator(q._ldf.visit(), engine).translate_ir() + child = Translator(q._ldf.visit(), streaming_engine).translate_ir() # Output (a_renamed, b) where a_renamed is Col("a") out_schema = {"a_renamed": child.schema["a"], "b": child.schema["b"]} exprs = ( @@ -437,9 +442,9 @@ def test_remap_partitioning_select_renamed_key(engine) -> None: assert result.local == "inherit" -def test_remap_partitioning_reorder_columns(engine) -> None: +def test_remap_partitioning_reorder_columns(streaming_engine) -> None: # Select (b, a) from (a, b, c) -> partition keys (a,b) become indices (1, 0) in output - select = _make_select_ir(engine, ("b", "a")) + select = _make_select_ir(streaming_engine, ("b", "a")) part = Partitioning(inter_rank=HashScheme((0, 1), 8), local="inherit") result = maybe_remap_partitioning(select, part) assert result is not None @@ -448,9 +453,9 @@ def test_remap_partitioning_reorder_columns(engine) -> None: assert result.inter_rank.modulus == 8 -def test_remap_partitioning_reorder_columns_projection(engine) -> None: +def test_remap_partitioning_reorder_columns_projection(streaming_engine) -> None: q = pl.LazyFrame({"a": [1], "b": [2], "c": [3]}) - child = Translator(q._ldf.visit(), engine).translate_ir() + child = Translator(q._ldf.visit(), streaming_engine).translate_ir() # Projection output (b, a) -> child has (a, b, c); partition keys (a,b) -> indices (1, 0) out_schema = {k: child.schema[k] for k in ("b", "a")} proj = Projection(out_schema, child) diff --git a/python/cudf_polars/tests/experimental/test_parallel.py b/python/cudf_polars/tests/experimental/test_parallel.py index 873c4623788e..181513b12f4d 100644 --- a/python/cudf_polars/tests/experimental/test_parallel.py +++ b/python/cudf_polars/tests/experimental/test_parallel.py @@ -20,7 +20,7 @@ @pytest.mark.parametrize("column", ["a", "b"]) -def test_explode_multi(column, engine): +def test_explode_multi(column, streaming_engine): df = pl.LazyFrame( { "a": [[1, 2], [3, 4], None], @@ -29,19 +29,19 @@ def test_explode_multi(column, engine): } ) q = df.explode(column) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) @pytest.mark.parametrize( "mapping", [{}, {"b": "c"}, {"b": "a", "a": "b"}, {"a": "c", "b": "d"}] ) -def test_rename_multi(mapping, engine): +def test_rename_multi(mapping, streaming_engine): df = pl.LazyFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) q = df.rename(mapping) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) -def test_rename_concat(engine) -> None: +def test_rename_concat(streaming_engine) -> None: # https://github.com/rapidsai/cudf/pull/19121#issuecomment-2959305678 q = pl.concat( [ @@ -49,10 +49,10 @@ def test_rename_concat(engine) -> None: pl.LazyFrame({"a": [4, 5, 6]}).rename({"a": "A"}), ] ) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) -def test_fallback_on_concat_zlice(engine) -> None: +def test_fallback_on_concat_zlice(streaming_engine) -> None: q = pl.concat( [ pl.LazyFrame({"a": [1, 2]}), @@ -64,7 +64,7 @@ def test_fallback_on_concat_zlice(engine) -> None: with pytest.raises( UserWarning, match="This slice not supported for multiple partitions." ): - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) # --------------------------------------------------------------------------- @@ -72,13 +72,13 @@ def test_fallback_on_concat_zlice(engine) -> None: # --------------------------------------------------------------------------- -def test_evaluate_streaming(engine): +def test_evaluate_streaming(streaming_engine): df = pl.LazyFrame({"a": [1, 2, 3], "b": [3, 4, 5], "c": [5, 6, 7], "d": [7, 9, 8]}) q = df.select(pl.col("a") - (pl.col("b") + pl.col("c") * 2), pl.col("d")).sort("d") expected = q.collect(engine="cpu") got_gpu = q.collect(engine=pl.GPUEngine(raise_on_fail=True)) - got_streaming = q.collect(engine=engine) + got_streaming = q.collect(engine=streaming_engine) assert_frame_equal(expected, got_gpu) assert_frame_equal(expected, got_streaming) @@ -132,7 +132,7 @@ def test_pickle_conditional_join_args(): @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -144,7 +144,7 @@ def test_pickle_conditional_join_args(): ], indirect=True, ) -def test_preserve_partitioning(engine): +def test_preserve_partitioning(streaming_engine): left = pl.LazyFrame({"a": [1, 2, 3, 4] * 5, "b": range(20)}) right = pl.LazyFrame({"a": [3, 4, 5, 6, 7] * 4, "c": range(20)}) q = ( @@ -169,4 +169,4 @@ def test_preserve_partitioning(engine): expect_dtype = ir.schema["a"] expect_expr = (NamedExpr("a", Col(expect_dtype, "a")),) assert partition_info[ir].partitioned_on == expect_expr - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) diff --git a/python/cudf_polars/tests/experimental/test_ray.py b/python/cudf_polars/tests/experimental/test_ray.py index 21725a040bf9..ec011c329228 100644 --- a/python/cudf_polars/tests/experimental/test_ray.py +++ b/python/cudf_polars/tests/experimental/test_ray.py @@ -35,9 +35,6 @@ def engine() -> Iterator[RayEngine]: pytestmark = [ - # Ray's internal subprocess management leaks /dev/null file handles; - # suppress the resulting ResourceWarning noise from its internals. - pytest.mark.filterwarnings("ignore::ResourceWarning"), pytest.mark.skipif( is_running_with_rrun(), reason="RayEngine must not be created from within an rrun cluster", diff --git a/python/cudf_polars/tests/experimental/test_scan.py b/python/cudf_polars/tests/experimental/test_scan.py index 18d7ae05bec5..6879c2882923 100644 --- a/python/cudf_polars/tests/experimental/test_scan.py +++ b/python/cudf_polars/tests/experimental/test_scan.py @@ -33,14 +33,14 @@ def df(): ("parquet", pl.scan_parquet), ], ) -def test_parallel_scan(tmp_path, df, fmt, scan_fn, engine): +def test_parallel_scan(tmp_path, df, fmt, scan_fn, streaming_engine): make_partitioned_source(df, tmp_path, fmt, n_files=3) q = scan_fn(tmp_path) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": {"target_partition_size": 1_000}, @@ -49,9 +49,9 @@ def test_parallel_scan(tmp_path, df, fmt, scan_fn, engine): ], indirect=True, ) -def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, engine): +def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine): make_partitioned_source(df, tmp_path, "parquet", n_files=1) - assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=engine) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) # --------------------------------------------------------------------------- @@ -60,44 +60,44 @@ def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, engine): @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 1_000}}], indirect=True, ) -def test_split_scan_aligns_to_row_group_boundaries(tmp_path, df, engine): +def test_split_scan_aligns_to_row_group_boundaries(tmp_path, df, streaming_engine): make_partitioned_source(df, tmp_path, "parquet", n_files=1, row_group_size=10) q = pl.scan_parquet(tmp_path) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) @pytest.mark.parametrize("mask", [None, pl.col("x") < 1_000]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [{"executor_options": {"target_partition_size": 1_000}}], indirect=True, ) -def test_split_scan_predicate(tmp_path, df, mask, engine): +def test_split_scan_predicate(tmp_path, df, mask, streaming_engine): make_partitioned_source(df, tmp_path, "parquet", n_files=1) q = pl.scan_parquet(tmp_path) if mask is not None: q = q.filter(mask) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) @pytest.mark.parametrize("n_files", [2, 3]) @pytest.mark.parametrize( - "blocksize,engine", + "blocksize,streaming_engine", [ (1_000, {"executor_options": {"target_partition_size": 1_000}}), (10_000, {"executor_options": {"target_partition_size": 10_000}}), (1_000_000, {"executor_options": {"target_partition_size": 1_000_000}}), ], - indirect=["engine"], + indirect=["streaming_engine"], ) -def test_target_partition_size(tmp_path, df, blocksize, n_files, engine): +def test_target_partition_size(tmp_path, df, blocksize, n_files, streaming_engine): make_partitioned_source(df, tmp_path, "parquet", n_files=n_files) q = pl.scan_parquet(tmp_path) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=streaming_engine) # Check partitioning (throwaway engine — no cluster/runtime needed) _engine = pl.GPUEngine( diff --git a/python/cudf_polars/tests/experimental/test_shuffler.py b/python/cudf_polars/tests/experimental/test_shuffler.py index 63a1c7bc9b25..00b1fdd7509f 100644 --- a/python/cudf_polars/tests/experimental/test_shuffler.py +++ b/python/cudf_polars/tests/experimental/test_shuffler.py @@ -19,7 +19,7 @@ @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -38,7 +38,7 @@ ], indirect=True, ) -def test_join_rapidsmpf(engine) -> None: +def test_join_rapidsmpf(streaming_engine) -> None: left = pl.LazyFrame( { "x": range(15), @@ -54,11 +54,11 @@ def test_join_rapidsmpf(engine) -> None: } ) q = left.join(right, on="y", how="inner") - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -75,7 +75,7 @@ def test_join_rapidsmpf(engine) -> None: ], indirect=True, ) -def test_sort_rapidsmpf(engine) -> None: +def test_sort_rapidsmpf(streaming_engine) -> None: df = pl.LazyFrame( { "x": range(15), @@ -84,7 +84,7 @@ def test_sort_rapidsmpf(engine) -> None: } ) q = df.sort(by=["y", "z"]) - assert_gpu_result_equal(q, engine=engine, check_row_order=True) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=True) def test_is_already_partitioned(): diff --git a/python/cudf_polars/tests/experimental/test_spilling.py b/python/cudf_polars/tests/experimental/test_spilling.py index 444b82c40320..cc2992944469 100644 --- a/python/cudf_polars/tests/experimental/test_spilling.py +++ b/python/cudf_polars/tests/experimental/test_spilling.py @@ -37,7 +37,7 @@ def create_test_table(nbytes: int, stream: Stream) -> plc.Table: @pytest.mark.parametrize( - "engine,spilled_host_mem_type", + "streaming_engine,spilled_host_mem_type", [ pytest.param( {"rapidsmpf_options": {"pinned_memory": "true"}}, @@ -49,19 +49,19 @@ def create_test_table(nbytes: int, stream: Stream) -> plc.Table: ), ({"rapidsmpf_options": {"pinned_memory": "false"}}, MemoryType.HOST), ], - indirect=["engine"], + indirect=["streaming_engine"], ) def test_make_spill_function( - engine: SPMDEngine, spilled_host_mem_type: MemoryType + streaming_engine: SPMDEngine, spilled_host_mem_type: MemoryType ) -> None: """Test that spilling prioritizes longest queues and newest messages.""" - context = engine.context + context = streaming_engine.context if spilled_host_mem_type == MemoryType.PINNED_HOST: - assert engine.context.br().pinned_mr is not None + assert streaming_engine.context.br().pinned_mr is not None other_host_mem_type = MemoryType.HOST else: - assert engine.context.br().pinned_mr is None + assert streaming_engine.context.br().pinned_mr is None other_host_mem_type = MemoryType.PINNED_HOST # Create 3 spillable message containers simulating fanout buffers diff --git a/python/cudf_polars/tests/experimental/test_statistics.py b/python/cudf_polars/tests/experimental/test_statistics.py index 686658de92e4..965449b80f04 100644 --- a/python/cudf_polars/tests/experimental/test_statistics.py +++ b/python/cudf_polars/tests/experimental/test_statistics.py @@ -24,9 +24,6 @@ # variants skip themselves in that environment. pytestmark = [ pytest.mark.spmd, - # Ray's subprocess management and distributed's shutdown leak unclosed - # /dev/null handles and sockets; suppress the noise. - pytest.mark.filterwarnings("ignore::ResourceWarning"), ] diff --git a/python/cudf_polars/tests/experimental/test_stats.py b/python/cudf_polars/tests/experimental/test_stats.py index 9fda8a32231f..50212be06717 100644 --- a/python/cudf_polars/tests/experimental/test_stats.py +++ b/python/cudf_polars/tests/experimental/test_stats.py @@ -151,7 +151,7 @@ def test_dataframe_round_trip() -> None: @pytest.mark.parametrize("kind", ["parquet", "csv", "frame"]) @pytest.mark.parametrize( - "engine", + "streaming_engine", [ { "executor_options": { @@ -162,7 +162,7 @@ def test_dataframe_round_trip() -> None: ], indirect=True, ) -def test_stats_planning(tmp_path, kind, engine): +def test_stats_planning(tmp_path, kind, streaming_engine): sales = pl.DataFrame( { "order_id": [1, 2, 3, 4, 5, 6], @@ -191,7 +191,7 @@ def test_stats_planning(tmp_path, kind, engine): pl.col("region").first().alias("region"), ] ) - assert_gpu_result_equal(q_gb.sort("customer_id"), engine=engine) + assert_gpu_result_equal(q_gb.sort("customer_id"), engine=streaming_engine) def test_parquet_deserialize_wrong_type() -> None: diff --git a/python/cudf_polars/tests/experimental/test_union.py b/python/cudf_polars/tests/experimental/test_union.py index cfadc59efc97..79a3ca649632 100644 --- a/python/cudf_polars/tests/experimental/test_union.py +++ b/python/cudf_polars/tests/experimental/test_union.py @@ -8,7 +8,7 @@ from cudf_polars.testing.asserts import assert_gpu_result_equal -def test_union_shared_fanout_no_deadlock(engine): +def test_union_shared_fanout_no_deadlock(streaming_engine): # union actor can deadlock when input branches share a fanout. # See https://github.com/rapidsai/cudf/issues/21750 n = 100 @@ -16,4 +16,4 @@ def test_union_shared_fanout_no_deadlock(engine): gb = df.group_by("key").agg(pl.col("val").sum()) project = df.select("key", "val") q = pl.concat([gb, project]) - assert_gpu_result_equal(q, engine=engine, check_row_order=False) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.py index 5ad6f68cc2ab..e465a87f0aa0 100644 --- a/python/cudf_polars/tests/expressions/test_agg.py +++ b/python/cudf_polars/tests/expressions/test_agg.py @@ -98,20 +98,20 @@ def decimal_df() -> pl.LazyFrame: ) -def test_agg(df, agg, xfail_if_sorted_gt_135): +def test_agg(engine: pl.GPUEngine, df, agg, xfail_if_sorted_gt_135): expr = getattr(pl.col("a"), agg)() q = df.select(expr) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) -def test_bool_agg(agg): +def test_bool_agg(engine: pl.GPUEngine, agg): if agg == "cum_min" or agg == "cum_max": pytest.skip("Does not apply") df = pl.LazyFrame({"a": [True, False, None, True]}) expr = getattr(pl.col("a"), agg)() q = df.select(expr) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) @pytest.mark.parametrize("cum_agg", sorted(expr.UnaryFunction._supported_cum_aggs)) @@ -125,10 +125,10 @@ def test_cum_agg_reverse_unsupported(cum_agg): @pytest.mark.parametrize("q", [0.5, pl.lit(0.5)]) @pytest.mark.parametrize("interp", ["nearest", "higher", "lower", "midpoint", "linear"]) -def test_quantile(df, q, interp, xfail_if_sorted_gt_135): +def test_quantile(engine: pl.GPUEngine, df, q, interp, xfail_if_sorted_gt_135): expr = pl.col("a").quantile(q, interp) q = df.select(expr) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) def test_quantile_invalid_q(df): @@ -152,7 +152,7 @@ def test_quantile_duration_unsupported(): @pytest.mark.parametrize( "op", [pl.Expr.min, pl.Expr.nan_min, pl.Expr.max, pl.Expr.nan_max] ) -def test_agg_float_with_nans(op): +def test_agg_float_with_nans(engine: pl.GPUEngine, op): df = pl.LazyFrame( { "a": pl.Series([1, 2, float("nan")], dtype=pl.Float64()), @@ -161,24 +161,24 @@ def test_agg_float_with_nans(op): ) q = df.select(op(pl.col("a")), op(pl.col("b"))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.xfail(reason="https://github.com/pola-rs/polars/issues/17513") @pytest.mark.parametrize("op", [pl.Expr.max, pl.Expr.min]) -def test_agg_singleton(op): +def test_agg_singleton(engine: pl.GPUEngine, op): df = pl.LazyFrame({"a": pl.Series([float("nan")])}) q = df.select(op(pl.col("a"))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("data", [[], [None], [None, 2, 3, None]]) -def test_sum_empty_zero(data): +def test_sum_empty_zero(engine: pl.GPUEngine, data): df = pl.LazyFrame({"a": pl.Series(values=data, dtype=pl.Int32())}) q = df.select(pl.col("a").sum()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_implode_agg_unsupported(): @@ -194,7 +194,7 @@ def test_implode_agg_unsupported(): assert_ir_translation_raises(q, NotImplementedError) -def test_decimal_aggs(decimal_df: pl.LazyFrame) -> None: +def test_decimal_aggs(engine: pl.GPUEngine, decimal_df: pl.LazyFrame) -> None: q = decimal_df.with_columns( sum=pl.col("a").sum(), min=pl.col("a").min(), @@ -204,25 +204,25 @@ def test_decimal_aggs(decimal_df: pl.LazyFrame) -> None: mean_f32=pl.col("a").mean().cast(pl.Float32), median_f32=pl.col("a").median().cast(pl.Float32), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("interp", ["nearest", "higher", "lower", "midpoint", "linear"]) -def test_decimal_quantile(decimal_df, interp): +def test_decimal_quantile(engine: pl.GPUEngine, decimal_df, interp): q = decimal_df.select(pl.col("a").quantile(0.5, interpolation=interp)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.skipif( POLARS_VERSION_LT_134, reason="std/var on decimal not supported before polars 1.34", ) -def test_decimal_std_var(decimal_df): +def test_decimal_std_var(engine: pl.GPUEngine, decimal_df): q = decimal_df.select( std=pl.col("a").std(), var=pl.col("a").var(), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_invalid_agg(request): @@ -237,10 +237,10 @@ def test_invalid_agg(request): assert_ir_translation_raises(q, NotImplementedError) -def test_sum_all_null_decimal_dtype(): +def test_sum_all_null_decimal_dtype(engine: pl.GPUEngine): df = pl.LazyFrame({"foo": pl.Series([None], dtype=pl.Decimal(9, 2))}) q = df.select(pl.col("foo").sum()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("expr", [pl.col("a").median(), pl.col("a").quantile(0.5)]) diff --git a/python/cudf_polars/tests/expressions/test_booleanfunction.py b/python/cudf_polars/tests/expressions/test_booleanfunction.py index 5344068e886a..3421d87df22e 100644 --- a/python/cudf_polars/tests/expressions/test_booleanfunction.py +++ b/python/cudf_polars/tests/expressions/test_booleanfunction.py @@ -28,7 +28,7 @@ def ignore_nulls(request: pytest.FixtureRequest) -> bool: return request.param -def test_booleanfunction_reduction(*, ignore_nulls: bool) -> None: +def test_booleanfunction_reduction(engine: pl.GPUEngine, *, ignore_nulls: bool) -> None: ldf = pl.LazyFrame( { "a": pl.Series([1, 2, 3.0, 2, 5], dtype=pl.Float64()), @@ -42,11 +42,11 @@ def test_booleanfunction_reduction(*, ignore_nulls: bool) -> None: (pl.col("b") > 2).all(ignore_nulls=ignore_nulls), ) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize("expr", [pl.Expr.any, pl.Expr.all]) -def test_booleanfunction_all_any_kleene(expr, ignore_nulls): +def test_booleanfunction_all_any_kleene(engine: pl.GPUEngine, expr, ignore_nulls): ldf = pl.LazyFrame( { "a": [False, None], @@ -61,7 +61,7 @@ def test_booleanfunction_all_any_kleene(expr, ignore_nulls): } ) q = ldf.select(expr(pl.col("*"), ignore_nulls=ignore_nulls)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -76,13 +76,14 @@ def test_booleanfunction_all_any_kleene(expr, ignore_nulls): ) @pytest.mark.parametrize("has_nans", [False, True], ids=["no_nans", "nans"]) def test_boolean_function_unary( + engine: pl.GPUEngine, expr: Callable[[pl.Expr], pl.Expr], *, has_nans: bool, has_nulls: bool, - using_rapidsmpf: bool, + using_streaming_engine: bool, ) -> None: - if using_rapidsmpf: + if using_streaming_engine: pytest.skip( "Avoiding possible segfault with cuda 12.9 builds https://github.com/rapidsai/cudf/issues/21828" ) @@ -96,7 +97,7 @@ def test_boolean_function_unary( q = df.select(expr(pl.col("a")), expr(pl.col("a")).not_().alias("b")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -107,7 +108,7 @@ def test_boolean_function_unary( pytest.param(lambda e: e.is_finite(), id="is_finite"), ], ) -def test_nan_in_non_floating_point_column(expr): +def test_nan_in_non_floating_point_column(engine: pl.GPUEngine, expr): ldf = pl.LazyFrame({"int": [-1, 1, None]}).with_columns( float=pl.col("int").cast(pl.Float64), float_na=pl.col("int") ** 0.5, @@ -121,7 +122,7 @@ def test_nan_in_non_floating_point_column(expr): ] ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -132,7 +133,7 @@ def test_nan_in_non_floating_point_column(expr): [pl.col("a").is_infinite(), pl.col("b").is_finite()], ], ) -def test_boolean_finite(expr): +def test_boolean_finite(engine: pl.GPUEngine, expr): df = pl.LazyFrame( { "a": pl.Series([1, float("nan"), 2, float("inf")], dtype=pl.Float64()), @@ -143,14 +144,14 @@ def test_boolean_finite(expr): q = df.select(expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("closed", ["both", "left", "right", "none"]) @pytest.mark.parametrize( "bounds", [(1, 2), (-1, 10), (11, 10), (pl.col("lo"), pl.col("hi"))] ) -def test_boolean_isbetween(closed, bounds): +def test_boolean_isbetween(engine: pl.GPUEngine, closed, bounds): df = pl.LazyFrame( { "a": pl.Series([1, float("nan"), 2, 4], dtype=pl.Float32()), @@ -159,17 +160,20 @@ def test_boolean_isbetween(closed, bounds): } ) - q = df.select(pl.col("a").is_between(*bounds, closed=closed)) + lower, upper = bounds + q = df.select(pl.col("a").is_between(lower, upper, closed=closed)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "expr", [pl.any_horizontal("*"), pl.all_horizontal("*")], ids=["any", "all"] ) @pytest.mark.parametrize("wide", [False, True], ids=["narrow", "wide"]) -def test_boolean_horizontal(expr, has_nulls, wide, using_rapidsmpf): - if using_rapidsmpf: +def test_boolean_horizontal( + engine: pl.GPUEngine, expr, has_nulls, wide, using_streaming_engine +): + if using_streaming_engine: pytest.skip( "Avoiding possible segfault with cuda 12.9 builds https://github.com/rapidsai/cudf/issues/21828" ) @@ -189,7 +193,7 @@ def test_boolean_horizontal(expr, has_nulls, wide, using_rapidsmpf): ldf = ldf.with_columns(pl.col("c").alias(f"col{i}") for i in range(128)) q = ldf.select(expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -198,14 +202,15 @@ def test_boolean_horizontal(expr, has_nulls, wide, using_rapidsmpf): pytest.param( pl.col("a").is_in(pl.col("b").implode()), marks=pytest.mark.xfail(reason="Need to support implode agg"), + id="implode", ), - pl.col("a").is_in([1, 2, 3]), - pl.col("a").is_in([]), - pl.col("a").is_in([3, 4, 2]), - pl.col("c").is_in([10, None, 11]), + pytest.param(pl.col("a").is_in([1, 2, 3]), id="list_small"), + pytest.param(pl.col("a").is_in([]), id="list_empty"), + pytest.param(pl.col("a").is_in([3, 4, 2]), id="list_shuffled"), + pytest.param(pl.col("c").is_in([10, None, 11]), id="list_with_nulls"), ], ) -def test_boolean_is_in(expr): +def test_boolean_is_in(engine: pl.GPUEngine, expr): ldf = pl.LazyFrame( { "a": pl.Series([1, 2, 3], dtype=pl.Int64()), @@ -217,11 +222,11 @@ def test_boolean_is_in(expr): q = ldf.select(expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("expr", [pl.Expr.and_, pl.Expr.or_, pl.Expr.xor]) -def test_boolean_kleene_logic(expr): +def test_boolean_kleene_logic(engine: pl.GPUEngine, expr): ldf = pl.LazyFrame( { "a": [False, False, False, None, None, None, True, True, True], @@ -229,7 +234,7 @@ def test_boolean_kleene_logic(expr): } ) q = ldf.select(expr(pl.col("a"), pl.col("b"))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_boolean_is_in_raises_unsupported(): @@ -240,17 +245,17 @@ def test_boolean_is_in_raises_unsupported(): assert_ir_translation_raises(q, NotImplementedError) -def test_boolean_is_in_with_nested_list_raises(): +def test_boolean_is_in_with_nested_list_raises(engine: pl.GPUEngine): ldf = pl.LazyFrame({"x": [1, 2, 3], "y": [[1, 2], [2, 3], [4]]}) q = ldf.select(pl.col("x").is_in(pl.col("y"))) with pytest.raises(AssertionError, match="DataFrames are different"): - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_expr_is_in_empty_list(): +def test_expr_is_in_empty_list(engine: pl.GPUEngine): ldf = pl.LazyFrame({"a": [1, 2, 3, 4]}) q = ldf.select(pl.col("a").is_in([])) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -285,7 +290,7 @@ def test_boolean_is_close(request): (pl.UInt8(), [1, 0, None, 255, 2]), ], ) -def test_boolean_not_with_integers(dtype, col): +def test_boolean_not_with_integers(engine: pl.GPUEngine, dtype, col): ldf = pl.LazyFrame({"a": pl.Series(col, dtype=dtype)}) q = ldf.select(~pl.col("a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_casting.py b/python/cudf_polars/tests/expressions/test_casting.py index 9fedf3e3658c..e8cd8c1fcddd 100644 --- a/python/cudf_polars/tests/expressions/test_casting.py +++ b/python/cudf_polars/tests/expressions/test_casting.py @@ -44,10 +44,10 @@ def tests(dtypes): @pytest.mark.parametrize("dtypes", _supported_dtypes, indirect=True) -def test_cast_supported(tests): +def test_cast_supported(engine: pl.GPUEngine, tests): df, totype = tests q = df.select(pl.col("a").cast(totype)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("dtypes", _unsupported_dtypes, indirect=True) @@ -58,15 +58,15 @@ def test_cast_unsupported(tests): ) -def test_allow_double_cast(): +def test_allow_double_cast(engine: pl.GPUEngine): df = pl.LazyFrame({"c0": [1000]}) query = df.select(pl.col("c0").cast(pl.Boolean).cast(pl.Int8)) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize("dtype", [pl.Int64(), pl.Float64()]) @pytest.mark.parametrize("strict", [True, False]) -def test_cast_strict_false_string_to_numeric(dtype, strict): +def test_cast_strict_false_string_to_numeric(engine: pl.GPUEngine, dtype, strict): df = pl.LazyFrame({"c0": ["1969-12-08 17:00:01", "1", None]}) query = df.with_columns(pl.col("c0").cast(dtype, strict=strict)) if strict: @@ -77,7 +77,7 @@ def test_cast_strict_false_string_to_numeric(dtype, strict): cudf_except=cudf_except, ) else: - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) def test_cast_from_string_unsupported(): @@ -92,7 +92,7 @@ def test_cast_to_string_unsupported(): assert_ir_translation_raises(query, NotImplementedError) -def test_float_to_decimal_rounding(): +def test_float_to_decimal_rounding(engine: pl.GPUEngine): # See https://github.com/rapidsai/cudf/pull/21450 df = pl.LazyFrame( { @@ -101,4 +101,4 @@ def test_float_to_decimal_rounding(): } ) q = df.select(pl.col("foo") / pl.col("bar")) - assert_gpu_result_equal(q, check_dtypes=not POLARS_VERSION_LT_132) + assert_gpu_result_equal(q, engine=engine, check_dtypes=not POLARS_VERSION_LT_132) diff --git a/python/cudf_polars/tests/expressions/test_datetime_basic.py b/python/cudf_polars/tests/expressions/test_datetime_basic.py index 78fa1c70a342..9f6fd36c5787 100644 --- a/python/cudf_polars/tests/expressions/test_datetime_basic.py +++ b/python/cudf_polars/tests/expressions/test_datetime_basic.py @@ -33,7 +33,7 @@ ], ids=repr, ) -def test_datetime_dataframe_scan(dtype): +def test_datetime_dataframe_scan(engine: pl.GPUEngine, dtype): ldf = pl.DataFrame( { "a": pl.Series([1, 2, 3, 4, 5, 6, 7], dtype=dtype), @@ -42,7 +42,7 @@ def test_datetime_dataframe_scan(dtype): ).lazy() query = ldf.select(pl.col("b"), pl.col("a")) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) datetime_extract_fields = [ @@ -67,7 +67,7 @@ def field(request): return request.param -def test_datetime_extract(field): +def test_datetime_extract(engine: pl.GPUEngine, field): ldf = pl.LazyFrame( { "datetimes": pl.datetime_range( @@ -81,7 +81,7 @@ def test_datetime_extract(field): q = ldf.select(field(pl.col("datetimes").dt)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_datetime_extra_unsupported(monkeypatch): @@ -122,7 +122,7 @@ def unsupported_name_getter(self): methodcaller("weekday"), ], ) -def test_date_extract(field): +def test_date_extract(engine: pl.GPUEngine, field): ldf = pl.LazyFrame( { "dates": [ @@ -138,11 +138,11 @@ def test_date_extract(field): q = ldf.select(field(pl.col("dates").dt)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("format", ["%Y-%m-%d", "%Y/%m/%d", "%Y.%m.%d", ""]) -def test_strftime_timestamp(format): +def test_strftime_timestamp(engine: pl.GPUEngine, format): ldf = pl.LazyFrame( { "dates": [ @@ -154,7 +154,7 @@ def test_strftime_timestamp(format): q = ldf.select(pl.col("dates").dt.strftime(format)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("format", ["iso", "polars"]) @@ -175,7 +175,7 @@ def test_strftime_duration(format): @pytest.mark.parametrize( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] ) -def test_datetime_month_start(dtype): +def test_datetime_month_start(engine: pl.GPUEngine, dtype): data = pl.DataFrame( { "dates": pl.Series( @@ -193,13 +193,13 @@ def test_datetime_month_start(dtype): ).lazy() q = data.select(pl.col("dates").dt.month_start()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] ) -def test_datetime_month_end(dtype): +def test_datetime_month_end(engine: pl.GPUEngine, dtype): data = pl.DataFrame( { "dates": pl.Series( @@ -217,7 +217,7 @@ def test_datetime_month_end(dtype): ).lazy() q = data.select(pl.col("dates").dt.month_end()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -234,11 +234,11 @@ def test_datetime_month_end(dtype): @pytest.mark.parametrize( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] ) -def test_is_leap_year(data, dtype): +def test_is_leap_year(engine: pl.GPUEngine, data, dtype): ldf = pl.LazyFrame({"dates": pl.Series(data, dtype=dtype)}) q = ldf.select(pl.col("dates").dt.is_leap_year()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -251,17 +251,17 @@ def test_is_leap_year(data, dtype): (datetime.date(2021, 1, 1), datetime.date(2021, 1, 2)), ], ) -def test_ordinal_day(start_date, end_date): +def test_ordinal_day(engine: pl.GPUEngine, start_date, end_date): df = pl.DataFrame({"date": pl.date_range(start_date, end_date, eager=True)}).lazy() q = df.with_columns( pl.col("date").dt.ordinal_day().alias("day_of_year"), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_isoweek(): +def test_isoweek(engine: pl.GPUEngine): df = pl.DataFrame( { "date": [ @@ -278,10 +278,10 @@ def test_isoweek(): q = df.with_columns(pl.col("date").dt.week().alias("isoweek")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_isoyear(): +def test_isoyear(engine: pl.GPUEngine): df = pl.DataFrame( { "date": [ @@ -299,14 +299,14 @@ def test_isoyear(): q = df.with_columns(pl.col("date").dt.iso_year().alias("isoyear")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] ) @pytest.mark.parametrize("time_unit", ["ms", "us", "ns"]) -def test_datetime_cast_time_unit_datetime(dtype, time_unit): +def test_datetime_cast_time_unit_datetime(engine: pl.GPUEngine, dtype, time_unit): sr = pl.Series( "date", [ @@ -322,14 +322,14 @@ def test_datetime_cast_time_unit_datetime(dtype, time_unit): q = df.select(pl.col("date").dt.cast_time_unit(time_unit).alias("time_unit_ms")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "dtype", [pl.Duration("ms"), pl.Duration("us"), pl.Duration("ns")] ) @pytest.mark.parametrize("time_unit", ["ms", "us", "ns"]) -def test_datetime_cast_time_unit_duration(dtype, time_unit): +def test_datetime_cast_time_unit_duration(engine: pl.GPUEngine, dtype, time_unit): sr = pl.Series( "date", [ @@ -344,7 +344,7 @@ def test_datetime_cast_time_unit_duration(dtype, time_unit): df = pl.DataFrame({"date": sr}).lazy() q = df.select(pl.col("date").dt.cast_time_unit(time_unit).alias("time_unit_ms")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -368,7 +368,7 @@ def test_datetime_cast_time_unit_duration(dtype, time_unit): pl.UInt8(), ], ) -def test_datetime_from_integer(datetime_dtype, integer_dtype): +def test_datetime_from_integer(engine: pl.GPUEngine, datetime_dtype, integer_dtype): values = [ 0, 1, @@ -385,7 +385,7 @@ def test_datetime_from_integer(datetime_dtype, integer_dtype): polars_except=pl.exceptions.InvalidOperationError, ) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -411,7 +411,7 @@ def test_datetime_from_integer(datetime_dtype, integer_dtype): pl.UInt8(), ], ) -def test_integer_from_datetime(datetime_dtype, integer_dtype): +def test_integer_from_datetime(engine: pl.GPUEngine, datetime_dtype, integer_dtype): values = [ 0, 1, @@ -421,4 +421,4 @@ def test_integer_from_datetime(datetime_dtype, integer_dtype): ] df = pl.LazyFrame({"data": pl.Series(values, dtype=datetime_dtype)}) q = df.select(pl.col("data").cast(integer_dtype).alias("int_from_datetime")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_distinct.py b/python/cudf_polars/tests/expressions/test_distinct.py index 9ab36bea3ef0..5812a23b321f 100644 --- a/python/cudf_polars/tests/expressions/test_distinct.py +++ b/python/cudf_polars/tests/expressions/test_distinct.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 @@ -25,7 +25,7 @@ def df(*, with_nulls: bool) -> pl.LazyFrame: return pl.LazyFrame({"a": values}) -def test_expr_distinct(df, op): +def test_expr_distinct(engine: pl.GPUEngine, df, op): expr = getattr(pl.col("a"), op)() query = df.select(expr) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_filter.py b/python/cudf_polars/tests/expressions/test_filter.py index 03a2db7d57a5..7b0cb00fcc60 100644 --- a/python/cudf_polars/tests/expressions/test_filter.py +++ b/python/cudf_polars/tests/expressions/test_filter.py @@ -21,7 +21,7 @@ ], ) @pytest.mark.parametrize("predicate_pushdown", [False, True]) -def test_filter_expression(expr, predicate_pushdown): +def test_filter_expression(engine: pl.GPUEngine, expr, predicate_pushdown): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -33,6 +33,7 @@ def test_filter_expression(expr, predicate_pushdown): query = ldf.select(pl.col("a").filter(expr)) assert_gpu_result_equal( query, + engine=engine, collect_kwargs={ "optimizations": pl.QueryOptFlags(predicate_pushdown=predicate_pushdown) }, diff --git a/python/cudf_polars/tests/expressions/test_gather.py b/python/cudf_polars/tests/expressions/test_gather.py index 402a0645b646..268c19025e6d 100644 --- a/python/cudf_polars/tests/expressions/test_gather.py +++ b/python/cudf_polars/tests/expressions/test_gather.py @@ -9,7 +9,7 @@ from cudf_polars.testing.asserts import assert_gpu_result_equal -def test_gather(): +def test_gather(engine: pl.GPUEngine): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -18,10 +18,10 @@ def test_gather(): ) query = ldf.select(pl.col("a").gather(pl.col("b"))) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_gather_with_nulls(): +def test_gather_with_nulls(engine: pl.GPUEngine): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -31,10 +31,10 @@ def test_gather_with_nulls(): query = ldf.select(pl.col("a").gather(pl.col("b"))) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_gather_empty_indices(): +def test_gather_empty_indices(engine: pl.GPUEngine): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5], @@ -42,11 +42,11 @@ def test_gather_empty_indices(): ) query = ldf.select(pl.col("a").gather(pl.lit(pl.Series("idx", [], dtype=pl.Int64)))) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize("negative", [False, True]) -def test_gather_out_of_bounds(negative): +def test_gather_out_of_bounds(engine_raise_on_fail: pl.GPUEngine, negative): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -57,7 +57,7 @@ def test_gather_out_of_bounds(negative): query = ldf.select(pl.col("a").gather(pl.col("b"))) with pytest.raises(ValueError, match="gather indices are out of bounds"): - query.collect(engine="gpu") + query.collect(engine=engine_raise_on_fail) @pytest.mark.parametrize( @@ -83,6 +83,7 @@ def test_gather_out_of_bounds(negative): ], ) def test_gather_on_literal( + engine: pl.GPUEngine, lit: pl.Expr, idx: pl.Expr, ) -> None: @@ -96,4 +97,4 @@ def test_gather_on_literal( ) q = df.select(lit.gather(idx)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_len.py b/python/cudf_polars/tests/expressions/test_len.py index ccc903aa99d8..d2d29818bf12 100644 --- a/python/cudf_polars/tests/expressions/test_len.py +++ b/python/cudf_polars/tests/expressions/test_len.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("dtype", [pl.UInt32, pl.Int32, None]) @pytest.mark.parametrize("empty", [False, True]) -def test_len(dtype, empty): +def test_len(engine: pl.GPUEngine, dtype, empty): if empty: df = pl.LazyFrame({}) else: @@ -25,16 +25,17 @@ def test_len(dtype, empty): # Workaround for https://github.com/pola-rs/polars/issues/16904 assert_gpu_result_equal( q, + engine=engine, collect_kwargs={"optimizations": pl.QueryOptFlags(projection_pushdown=False)}, ) @pytest.mark.parametrize("data", [[1, 2, 3], [1, 2, None]]) -def test_col_len(data): +def test_col_len(engine: pl.GPUEngine, data): data = {"a": list("xyz"), "b": data} q = pl.LazyFrame(data).select( pl.col("a").len().alias("l"), (pl.col("a").len() * 2).alias("l2"), pl.col("b").len().alias("l3"), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_literal.py b/python/cudf_polars/tests/expressions/test_literal.py index 2a32d92a6203..77e95ce26624 100644 --- a/python/cudf_polars/tests/expressions/test_literal.py +++ b/python/cudf_polars/tests/expressions/test_literal.py @@ -40,12 +40,12 @@ def float(request): return pl.lit(1.0, dtype=request.param) -def test_numeric_literal(integer, float): +def test_numeric_literal(engine: pl.GPUEngine, integer, float): df = pl.LazyFrame({}) q = df.select(integer=integer, float_=float, sum_=integer + float) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.fixture( @@ -60,7 +60,7 @@ def timedelta(request): return pl.lit(9_000, dtype=request.param) -def test_timelike_literal(timestamp, timedelta): +def test_timelike_literal(engine: pl.GPUEngine, timestamp, timedelta): df = pl.LazyFrame({}) q = df.select( @@ -81,12 +81,12 @@ def test_timelike_literal(timestamp, timedelta): schema["delta"], plc.binaryop.BinaryOperator.ADD, ): - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) else: assert_ir_translation_raises(q, NotImplementedError) -def test_select_literal_series(): +def test_select_literal_series(engine: pl.GPUEngine): df = pl.LazyFrame({}) q = df.select( @@ -95,7 +95,7 @@ def test_select_literal_series(): c=pl.Series([[[1]], [], [[1, 2, 3, 4]]], dtype=pl.List(pl.List(pl.Float32()))), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( diff --git a/python/cudf_polars/tests/expressions/test_numeric_binops.py b/python/cudf_polars/tests/expressions/test_numeric_binops.py index 82c64e772616..32c203423dc3 100644 --- a/python/cudf_polars/tests/expressions/test_numeric_binops.py +++ b/python/cudf_polars/tests/expressions/test_numeric_binops.py @@ -72,46 +72,46 @@ def df(request, ltype, rtype, with_nulls, binop): return pl.LazyFrame({"a": a, "b": b}, schema={"a": ltype, "b": rtype}) -def test_numeric_binop(df, binop): +def test_numeric_binop(engine: pl.GPUEngine, df, binop): left = pl.col("a") right = pl.col("b") q = df.select(binop(left, right)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("left_scalar", [False, True]) @pytest.mark.parametrize("right_scalar", [False, True]) -def test_binop_with_scalar(left_scalar, right_scalar): +def test_binop_with_scalar(engine: pl.GPUEngine, left_scalar, right_scalar): df = pl.LazyFrame({"a": [1, 2, 3], "b": [5, 6, 7]}) lop = pl.lit(2) if left_scalar else pl.col("a") rop = pl.lit(6) if right_scalar else pl.col("b") q = df.select(lop / rop) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("zero", [0, pl.lit(0)]) -def test_floor_div_binop_by_zero(zero, ltype): +def test_floor_div_binop_by_zero(engine: pl.GPUEngine, zero, ltype): df = pl.LazyFrame({"a": [1, 0, 3]}, schema={"a": ltype}) q = df.select(pl.col("a") // zero) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("divisor", [1, 2.0]) -def test_true_div_boolean_column(divisor): +def test_true_div_boolean_column(engine: pl.GPUEngine, divisor): df = pl.LazyFrame({"a": [True, False]}) q = df.select(pl.col("a") / divisor) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_true_div_with_decimals(): +def test_true_div_with_decimals(engine: pl.GPUEngine): df = pl.LazyFrame( { "foo": [Decimal("1.00"), Decimal("2.00"), Decimal("3.00"), None], @@ -120,10 +120,10 @@ def test_true_div_with_decimals(): schema={"foo": pl.Decimal(15, 2), "bar": pl.Decimal(15, 2)}, ) q = df.select(pl.col("bar") / pl.col("foo")) - assert_gpu_result_equal(q, check_dtypes=not POLARS_VERSION_LT_132) + assert_gpu_result_equal(q, engine=engine, check_dtypes=not POLARS_VERSION_LT_132) -def test_multiply_with_decimals(): +def test_multiply_with_decimals(engine: pl.GPUEngine): df = pl.LazyFrame( { "x": [Decimal("1.23"), Decimal("4.56"), Decimal("7.89")], @@ -133,7 +133,7 @@ def test_multiply_with_decimals(): ) q = df.select(pl.col("x") * pl.col("y")) - assert_gpu_result_equal(q, check_dtypes=not POLARS_VERSION_LT_132) + assert_gpu_result_equal(q, engine=engine, check_dtypes=not POLARS_VERSION_LT_132) def test_sum_decimal_widens_precision(request) -> None: diff --git a/python/cudf_polars/tests/expressions/test_numeric_unaryops.py b/python/cudf_polars/tests/expressions/test_numeric_unaryops.py index 3cfbf0906c93..8286c5340fef 100644 --- a/python/cudf_polars/tests/expressions/test_numeric_unaryops.py +++ b/python/cudf_polars/tests/expressions/test_numeric_unaryops.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 @@ -71,25 +71,25 @@ def ldf(with_nulls, dtype): ) -def test_unary(ldf, op): +def test_unary(engine: pl.GPUEngine, ldf, op): expr = getattr(pl.col("a"), op)() q = ldf.select(expr) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) @pytest.mark.parametrize("base_literal", [False, True]) @pytest.mark.parametrize("exponent_literal", [False, True]) -def test_pow(ldf, base_literal, exponent_literal): +def test_pow(engine: pl.GPUEngine, ldf, base_literal, exponent_literal): base = pl.lit(2) if base_literal else pl.col("a") exponent = pl.lit(-3, dtype=pl.Float32) if exponent_literal else pl.col("b") q = ldf.select(base.pow(exponent)) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) @pytest.mark.parametrize("natural", [True, False]) -def test_log(ldf, natural): +def test_log(engine: pl.GPUEngine, ldf, natural): if natural: expr = pl.col("a").log() else: @@ -97,16 +97,16 @@ def test_log(ldf, natural): q = ldf.select(expr) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) @pytest.mark.parametrize("col", ["a", "b", "c"]) -def test_negate(ldf, col): +def test_negate(engine: pl.GPUEngine, ldf, col): q = ldf.select(-pl.col(col)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_null_count(): +def test_null_count(engine: pl.GPUEngine): lf = pl.LazyFrame( { "foo": [1, None, 3], @@ -119,27 +119,38 @@ def test_null_count(): pl.col("bar").is_null().sum(), pl.col("baz").is_null().sum(), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) def test_rank_supported( - request, ldf: pl.LazyFrame, method: RankMethod, *, descending: bool + engine: pl.GPUEngine, + request, + ldf: pl.LazyFrame, + method: RankMethod, + *, + descending: bool, ): request.applymarker( pytest.mark.xfail(condition=POLARS_VERSION_LT_132, reason="rank unsupported") ) expr = pl.col("a").rank(method=method, descending=descending) q = ldf.select(expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("test", ["with_nulls", "with_ties"]) def test_rank_methods_with_nulls_or_ties( - request, ldf: pl.LazyFrame, method: RankMethod, *, descending: bool, test: str + engine: pl.GPUEngine, + request, + ldf: pl.LazyFrame, + method: RankMethod, + *, + descending: bool, + test: str, ) -> None: request.applymarker( pytest.mark.xfail(condition=POLARS_VERSION_LT_132, reason="rank unsupported") @@ -152,7 +163,7 @@ def test_rank_methods_with_nulls_or_ties( expr = pl.when((base % 2) == 0).then(pl.lit(-5)).otherwise(base) q = ldf.select(expr.rank(method=method, descending=descending)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("seed", [42]) @@ -164,6 +175,6 @@ def test_rank_unsupported(ldf: pl.LazyFrame, method: RankMethod, seed: int) -> N @pytest.mark.parametrize("mode", ["half_to_even", "half_away_from_zero"]) -def test_round(ldf: pl.LazyFrame, mode: RoundMethod) -> None: +def test_round(engine: pl.GPUEngine, ldf: pl.LazyFrame, mode: RoundMethod) -> None: q = ldf.select(pl.col("a").sin().round(2, mode=mode)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index 4b40178e8bfd..776196753b68 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -33,7 +33,7 @@ def df(): @pytest.mark.parametrize("time_unit", ["ns", "us", "ms"]) -def test_rolling_datetime(request, time_unit): +def test_rolling_datetime(engine: pl.GPUEngine, request, time_unit): request.applymarker( pytest.mark.xfail( condition=not POLARS_VERSION_LT_136, @@ -59,10 +59,10 @@ def test_rolling_datetime(request, time_unit): max_a=pl.max("a").rolling(index_column="dt", period="10d", offset="2d"), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_rolling_date(request): +def test_rolling_date(engine: pl.GPUEngine, request): if not POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") @@ -84,11 +84,11 @@ def test_rolling_date(request): max_a=pl.max("a").rolling(index_column="dt", period="10d", offset="2d"), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("dtype", [pl.Int32, pl.UInt32, pl.Int64, pl.UInt64]) -def test_rolling_integral_orderby(request, dtype): +def test_rolling_integral_orderby(engine: pl.GPUEngine, request, dtype): if not POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") @@ -103,7 +103,7 @@ def test_rolling_integral_orderby(request, dtype): pl.col("values").sum().rolling("orderby", period="4i", closed="both") ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.skipif( @@ -129,7 +129,7 @@ def test_rolling_collect_list_raises(): ) -def test_unsorted_raises(request): +def test_unsorted_raises(engine_raise_on_fail: pl.GPUEngine, request): if not POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") @@ -142,10 +142,10 @@ def test_unsorted_raises(request): RuntimeError, match=r"Index column.*in rolling is not sorted, please sort first", ): - q.collect(engine=pl.GPUEngine(raise_on_fail=True)) + q.collect(engine=engine_raise_on_fail) -def test_orderby_nulls_raises_computeerror(request): +def test_orderby_nulls_raises_computeerror(engine_raise_on_fail: pl.GPUEngine, request): if not POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") @@ -157,7 +157,7 @@ def test_orderby_nulls_raises_computeerror(request): with pytest.raises( RuntimeError, match=r"Index column.*in rolling may not contain nulls" ): - q.collect(engine=pl.GPUEngine(raise_on_fail=True)) + q.collect(engine=engine_raise_on_fail) def test_invalid_duration_spec_raises_in_translation(request): @@ -188,7 +188,7 @@ def test_rolling_inside_groupby_raises(request): assert_ir_translation_raises(q, NotImplementedError) -def test_rolling_sum_all_null_window_returns_null(request): +def test_rolling_sum_all_null_window_returns_null(engine: pl.GPUEngine, request): if not POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") @@ -203,7 +203,7 @@ def test_rolling_sum_all_null_window_returns_null(request): out=pl.col("null_windows").sum().rolling("orderby", period="2i", closed="both") ) # Expected: [null, null, 5, 5, 5, 1] - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -231,19 +231,19 @@ def test_rolling_sum_all_null_window_returns_null(request): "literal_partition", ], ) -def test_over_group_various(df, expr): +def test_over_group_various(engine: pl.GPUEngine, df, expr): q = df.select(expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_window_over_group_sum_all_null_group_is_zero(df): +def test_window_over_group_sum_all_null_group_is_zero(engine: pl.GPUEngine, df): q = df.with_columns( pl.when(pl.col("g") == 1) .then(pl.lit(None, dtype=pl.Int64)) .otherwise(pl.col("x")) .alias("null") ).select(s=pl.col("null").sum().over("g")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -259,7 +259,9 @@ def test_window_over_group_sum_all_null_group_is_zero(df): ) @pytest.mark.parametrize("order_by_descending", [False, True]) @pytest.mark.parametrize("order_by_nulls_last", [False, True]) -def test_over_with_order_by(df, order_by, order_by_descending, order_by_nulls_last): +def test_over_with_order_by( + engine: pl.GPUEngine, df, order_by, order_by_descending, order_by_nulls_last +): q = df.select( pl.col("x") .sum() @@ -270,7 +272,7 @@ def test_over_with_order_by(df, order_by, order_by_descending, order_by_nulls_la nulls_last=order_by_nulls_last, ) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("strategy", ["explode", "join"], ids=["explode", "join"]) @@ -284,7 +286,7 @@ def test_over_boolean_function_unsupported(df): assert_ir_translation_raises(q, NotImplementedError) -def test_over_ternary(df): +def test_over_ternary(engine: pl.GPUEngine, df): q = df.select( pl.when(pl.col("g") == 1) .then(pl.lit(None, dtype=pl.Int64)) @@ -293,10 +295,13 @@ def test_over_ternary(df): .over("g") ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_over_broadcast_input_row_group_indices_aligned(): +@pytest.mark.skip_on_streaming_engine( + "GroupedWindow not supported for multiple partitions" +) +def test_over_broadcast_input_row_group_indices_aligned(engine: pl.GPUEngine): num_rows, num_groups = 512, 64 df = pl.LazyFrame( @@ -307,13 +312,14 @@ def test_over_broadcast_input_row_group_indices_aligned(): ) q = df.select(pl.col("x").sum().over("g")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("order_by", [None, ["g2", pl.col("x2") * 2]]) def test_rank_over( + engine: pl.GPUEngine, request, df: pl.LazyFrame, method: RankMethod, @@ -329,13 +335,14 @@ def test_rank_over( .rank(method=method, descending=descending) .over("g", order_by=order_by) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("order_by", [None, ["g2", pl.col("x2") * 2]]) def test_rank_over_with_ties( + engine: pl.GPUEngine, request, df: pl.LazyFrame, method: RankMethod, @@ -353,13 +360,14 @@ def test_rank_over_with_ties( .rank(method=method, descending=descending) .over("g", order_by=order_by) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("order_by", [None, ["g2", pl.col("x2") * 2]]) def test_rank_over_with_null_values( + engine: pl.GPUEngine, request, df: pl.LazyFrame, method: RankMethod, @@ -377,13 +385,14 @@ def test_rank_over_with_null_values( .rank(method=method, descending=descending) .over("g", order_by=order_by) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("method", ["ordinal", "dense", "min", "max", "average"]) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("order_by", [None, ["g2", pl.col("x2") * 2]]) def test_rank_over_with_null_group_keys( + engine: pl.GPUEngine, request, df: pl.LazyFrame, method: RankMethod, @@ -399,7 +408,7 @@ def test_rank_over_with_null_group_keys( .rank(method=method, descending=descending) .over("g_null", order_by=order_by) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("strategy", ["forward", "backward"]) @@ -420,6 +429,7 @@ def test_rank_over_with_null_group_keys( ], ) def test_fill_over( + engine: pl.GPUEngine, df: pl.LazyFrame, strategy: str, order_by: None | list[str | pl.Expr], @@ -434,7 +444,7 @@ def test_fill_over( if POLARS_VERSION_LT_132: assert_ir_translation_raises(q, NotImplementedError) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_fill_null_with_mean_over_unsupported(df: pl.LazyFrame) -> None: @@ -458,6 +468,7 @@ def test_fill_null_with_mean_over_unsupported(df: pl.LazyFrame) -> None: ], ) def test_cum_sum_over( + engine: pl.GPUEngine, df: pl.LazyFrame, *, expr: pl.Expr, @@ -465,7 +476,7 @@ def test_cum_sum_over( order_by: None | list[str | pl.Expr], ) -> None: q = df.select(expr.cum_sum().over(group_key, order_by=order_by)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -484,7 +495,9 @@ def test_cum_sum_over( ["g_null", "g2"], ], ) -def test_order_sensitive_over_scalar_aggs(df, expr, descending, nulls_last, order_by): +def test_order_sensitive_over_scalar_aggs( + engine: pl.GPUEngine, df, expr, descending, nulls_last, order_by +): q = df.select( expr.over( "g", @@ -496,4 +509,4 @@ def test_order_sensitive_over_scalar_aggs(df, expr, descending, nulls_last, orde if isinstance(order_by, list): assert_ir_translation_raises(q, NotImplementedError) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_round.py b/python/cudf_polars/tests/expressions/test_round.py index 3af3a0ce6d18..e60338792e88 100644 --- a/python/cudf_polars/tests/expressions/test_round.py +++ b/python/cudf_polars/tests/expressions/test_round.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -26,7 +26,7 @@ def df(dtype, with_nulls): @pytest.mark.parametrize("decimals", [0, 2, 4]) -def test_round(df, decimals): +def test_round(engine: pl.GPUEngine, df, decimals): q = df.select(pl.col("a").round(decimals=decimals)) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) diff --git a/python/cudf_polars/tests/expressions/test_shift.py b/python/cudf_polars/tests/expressions/test_shift.py index dc7275c80508..ebd955c9f17a 100644 --- a/python/cudf_polars/tests/expressions/test_shift.py +++ b/python/cudf_polars/tests/expressions/test_shift.py @@ -12,10 +12,10 @@ @pytest.mark.parametrize("n", [0, 1, 2, -1, -2, 5, -5]) -def test_shift(n): +def test_shift(engine: pl.GPUEngine, n): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]}) q = df.select(pl.col("a").shift(n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -28,50 +28,50 @@ def test_shift(n): (0, 7), ], ) -def test_shift_and_fill(n, fill_value): +def test_shift_and_fill(engine: pl.GPUEngine, n, fill_value): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) q = df.select(pl.col("a").shift(n, fill_value=fill_value)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_multiple_columns(): +def test_shift_multiple_columns(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]}) q = df.select(pl.col("a").shift(2), pl.col("b").shift(-1)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_float(): +def test_shift_float(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) q = df.select(pl.col("a").shift(2)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_and_fill_float(): +def test_shift_and_fill_float(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) q = df.select(pl.col("a").shift(2, fill_value=0.0)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_and_fill_expr(): +def test_shift_and_fill_expr(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) q = df.select(pl.col("a").shift(n=2, fill_value=pl.col("a").min())) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_string(): +def test_shift_string(engine: pl.GPUEngine): df = pl.LazyFrame({"a": ["x", "y", "z"]}) q = df.select(pl.col("a").shift(1)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_with_columns(): +def test_shift_with_columns(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) q = df.with_columns(shift=pl.col("a").shift(-2, fill_value=100)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("n", [1, -1, 2]) -def test_shift_datetime(n): +def test_shift_datetime(engine: pl.GPUEngine, n): df = pl.LazyFrame( { "a": [ @@ -87,11 +87,11 @@ def test_shift_datetime(n): } ) q = df.select(pl.col("a").shift(n), pl.col("b").shift(-n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("n", [1, -1]) -def test_shift_date(n): +def test_shift_date(engine: pl.GPUEngine, n): df = pl.LazyFrame( { "a": [ @@ -102,22 +102,22 @@ def test_shift_date(n): } ) q = df.select(pl.col("a").shift(n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_by_expression(): +def test_shift_by_expression(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5], "b": [1, 1, 1, 1, 1]}) q = df.select(pl.col("a").shift(n=pl.col("b").first())) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_by_expression_last(): +def test_shift_by_expression_last(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5], "b": [2, 2, 2, 2, 2]}) q = df.select(pl.col("a").shift(n=pl.col("b").last())) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_shift_by_expression_get(): +def test_shift_by_expression_get(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3, 4, 5], "b": [2, 2, 2, 2, 2]}) q = df.select(pl.col("a").shift(n=pl.col("b").get(2))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_slice.py b/python/cudf_polars/tests/expressions/test_slice.py index 9873be2455fc..c8f3b5d45021 100644 --- a/python/cudf_polars/tests/expressions/test_slice.py +++ b/python/cudf_polars/tests/expressions/test_slice.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -17,8 +17,8 @@ (-1,), ], ) -def test_slice(zlice): +def test_slice(engine: pl.GPUEngine, zlice): df = pl.LazyFrame({"a": [0, 1, 2, 3], "b": [1, 2, 3, 4]}) q = df.select(pl.col("a").slice(*zlice)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index f0688c84daa6..df66b324db52 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -18,7 +18,7 @@ @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("nulls_last", [False, True]) -def test_sort_expression(descending, nulls_last): +def test_sort_expression(engine: pl.GPUEngine, descending, nulls_last): ldf = pl.LazyFrame( { "a": [5, -1, 3, 4, None, 8, 6, 7, None], @@ -26,7 +26,7 @@ def test_sort_expression(descending, nulls_last): ) query = ldf.select(pl.col("a").sort(descending=descending, nulls_last=nulls_last)) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize( @@ -36,7 +36,9 @@ def test_sort_expression(descending, nulls_last): "nulls_last", itertools.combinations_with_replacement([False, True], 3) ) @pytest.mark.parametrize("maintain_order", [False, True], ids=["unstable", "stable"]) -def test_sort_by_expression(descending, nulls_last, maintain_order): +def test_sort_by_expression( + engine: pl.GPUEngine, descending, nulls_last, maintain_order +): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], @@ -55,13 +57,13 @@ def test_sort_by_expression(descending, nulls_last, maintain_order): maintain_order=maintain_order, ) ) - assert_gpu_result_equal(query, check_row_order=maintain_order) + assert_gpu_result_equal(query, engine=engine, check_row_order=maintain_order) @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("nulls_last", [False, True]) @pytest.mark.parametrize("with_nulls", ["no_nulls", "nulls"]) -def test_setsorted(request, descending, nulls_last, with_nulls): +def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_nulls): if not POLARS_VERSION_LT_135 and POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail( @@ -69,14 +71,15 @@ def test_setsorted(request, descending, nulls_last, with_nulls): "fixed in https://github.com/pola-rs/polars/pull/25250" ) ) - values = sorted([1, 2, 3, 4, 5, 6, -2], reverse=descending) + sorted_values = sorted([1, 2, 3, 4, 5, 6, -2], reverse=descending) + values: list[int | None] = [*sorted_values] if with_nulls == "nulls": values[-1 if nulls_last else 0] = None - df = pl.LazyFrame({"a": values}) + ldf = pl.LazyFrame({"a": values}) - q = df.set_sorted("a", descending=descending) + q = ldf.set_sorted("a", descending=descending) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) if POLARS_VERSION_LT_135: translator = Translator(q._ldf.visit(), pl.GPUEngine()) @@ -101,7 +104,7 @@ def test_setsorted(request, descending, nulls_last, with_nulls): ) -def test_sort_concat_filtered_to_empty(): +def test_sort_concat_filtered_to_empty(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) q = pl.concat([df.filter(pl.col("a") == 0), df.filter(pl.col("a") == 4)]).sort("a") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_stringfunction.py b/python/cudf_polars/tests/expressions/test_stringfunction.py index 6b3d204b60a2..be4f876afcf9 100644 --- a/python/cudf_polars/tests/expressions/test_stringfunction.py +++ b/python/cudf_polars/tests/expressions/test_stringfunction.py @@ -144,14 +144,14 @@ def slice_column_data(ldf, request): return ldf.with_columns(pl.lit(start).alias("start")) -def test_supported_stringfunction_expression(ldf): +def test_supported_stringfunction_expression(engine: pl.GPUEngine, ldf): q = ldf.select( pl.col("a").str.starts_with("Z"), pl.col("a").str.ends_with("h").alias("endswith_h"), pl.col("a").str.to_lowercase().alias("lower"), pl.col("a").str.to_uppercase().alias("upper"), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_unsupported_stringfunction(ldf): @@ -186,22 +186,22 @@ def test_contains_re_non_literal_raises(ldf): "j|u", ], ) -def test_contains_regex(ldf, substr): +def test_contains_regex(engine: pl.GPUEngine, ldf, substr): q = ldf.select(pl.col("a").str.contains(substr)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "literal", ["A", "de", "FGHI", "j", "kLm", "nOPq", "RsT", "uVw"] ) -def test_contains_literal(ldf, literal): +def test_contains_literal(engine: pl.GPUEngine, ldf, literal): q = ldf.select(pl.col("a").str.contains(pl.lit(literal), literal=True)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_contains_column(ldf): +def test_contains_column(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.contains(pl.col("a"), literal=True)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_contains_invalid(ldf): @@ -214,15 +214,15 @@ def test_contains_invalid(ldf): @pytest.mark.parametrize("offset", [1, -1, 0, 100, -100]) -def test_slice_scalars_offset(ldf, offset): +def test_slice_scalars_offset(engine: pl.GPUEngine, ldf, offset): q = ldf.select(pl.col("a").str.slice(offset)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("offset,length", slice_cases) -def test_slice_scalars_length_and_offset(ldf, offset, length): +def test_slice_scalars_length_and_offset(engine: pl.GPUEngine, ldf, offset, length): q = ldf.select(pl.col("a").str.slice(offset, length)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_slice_column(slice_column_data): @@ -242,15 +242,15 @@ def ldf_split(): @pytest.mark.parametrize("n", [1, 2, 10]) @pytest.mark.parametrize("by", ["_", " "]) -def test_split_n(ldf_split, n, by): +def test_split_n(engine: pl.GPUEngine, ldf_split, n, by): q = ldf_split.select(pl.col("a").str.splitn(by, n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("n", [1, 2, 10]) -def test_split_exact(ldf_split, n): +def test_split_exact(engine: pl.GPUEngine, ldf_split, n): q = ldf_split.select(pl.col("a").str.split_exact("_", n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_split_exact_inclusive_unsupported(ldf_split): @@ -258,10 +258,10 @@ def test_split_exact_inclusive_unsupported(ldf_split): assert_ir_translation_raises(q, NotImplementedError) -def test_split_exact_null_correct_children(): +def test_split_exact_null_correct_children(engine: pl.GPUEngine): df = pl.LazyFrame({"a": ["a_b", None]}) q = df.slice(1).select(pl.col("a").str.split_exact("_", 1)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("cache", [True, False], ids=lambda cache: f"{cache=}") @@ -280,7 +280,9 @@ def test_split_exact_null_correct_children(): ], ids=["valid", "valid", "invalid", "invalid", "invalid", "valid"], ) -def test_to_datetime(values, has_invalid_row, cache, strict, format, exact): +def test_to_datetime( + engine: pl.GPUEngine, values, has_invalid_row, cache, strict, format, exact +): df = pl.DataFrame({"a": pl.Series(values, dtype=pl.String())}) q = df.lazy().select( pl.col("a").str.strptime( @@ -310,7 +312,7 @@ def test_to_datetime(values, has_invalid_row, cache, strict, format, exact): cudf_except=cudf_exc, ) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -318,9 +320,9 @@ def test_to_datetime(values, has_invalid_row, cache, strict, format, exact): [("a", "a"), ("Wı", "☺"), ("FG", ""), ("doesnotexist", "blahblah")], # noqa: RUF001 ) @pytest.mark.parametrize("n", [0, 3, -1]) -def test_replace_literal(ldf, target, repl, n): +def test_replace_literal(engine: pl.GPUEngine, ldf, target, repl, n): q = ldf.select(pl.col("a").str.replace(target, repl, literal=True, n=n)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("target, repl", [("", ""), ("a", pl.col("a"))]) @@ -347,16 +349,16 @@ def test_replace_re(ldf): ), ], ) -def test_replace_many(ldf, target, repl): +def test_replace_many(engine: pl.GPUEngine, ldf, target, repl): q = ldf.select(pl.col("a").str.replace_many(target, repl)) _need_support_for_implode_agg = isinstance(repl, list) if _need_support_for_implode_agg: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) elif POLARS_VERSION_LT_131: assert_ir_translation_raises(q, NotImplementedError) else: # Polars 1.31 now gives us replacement argument as a list - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -435,19 +437,19 @@ def to_strip(request): return request.param -def test_strip_chars(strip_ldf, to_strip): +def test_strip_chars(engine: pl.GPUEngine, strip_ldf, to_strip): q = strip_ldf.select(pl.col("a").str.strip_chars(to_strip)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_strip_chars_start(strip_ldf, to_strip): +def test_strip_chars_start(engine: pl.GPUEngine, strip_ldf, to_strip): q = strip_ldf.select(pl.col("a").str.strip_chars_start(to_strip)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_strip_chars_end(strip_ldf, to_strip): +def test_strip_chars_end(engine: pl.GPUEngine, strip_ldf, to_strip): q = strip_ldf.select(pl.col("a").str.strip_chars_end(to_strip)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_strip_chars_column(strip_ldf): @@ -478,22 +480,22 @@ def test_unsupported_regex_raises(pattern): assert_ir_translation_raises(q, NotImplementedError) -def test_string_to_integer(str_to_integer_data, integer_type): +def test_string_to_integer(engine: pl.GPUEngine, str_to_integer_data, integer_type): q = str_to_integer_data.select(pl.col("a").cast(integer_type)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_from_integer(str_from_integer_data): +def test_string_from_integer(engine: pl.GPUEngine, str_from_integer_data): q = str_from_integer_data.select(pl.col("a").cast(pl.String)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_to_float(str_to_float_data, floating_type): +def test_string_to_float(engine: pl.GPUEngine, str_to_float_data, floating_type): q = str_to_float_data.select(pl.col("a").cast(floating_type)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_from_float(request, str_from_float_data): +def test_string_from_float(engine: pl.GPUEngine, request, str_from_float_data): if str_from_float_data.collect_schema()["a"] == pl.Float32: # libcudf will return a string representing the precision out to # a certain number of hardcoded decimal places. This results in @@ -509,7 +511,7 @@ def test_string_from_float(request, str_from_float_data): # libcudf reads float('inf') -> "inf" # but polars reads float('inf') -> "Inf" q = q.select(pl.col("a").str.to_lowercase()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_string_to_numeric_invalid(numeric_type): @@ -524,9 +526,9 @@ def test_string_to_numeric_invalid(numeric_type): @pytest.mark.parametrize("ignore_nulls", [False, True]) @pytest.mark.parametrize("delimiter", ["", "/"]) -def test_string_join(ldf, ignore_nulls, delimiter): +def test_string_join(engine: pl.GPUEngine, ldf, ignore_nulls, delimiter): q = ldf.select(pl.col("a").str.join(delimiter, ignore_nulls=ignore_nulls)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -550,7 +552,7 @@ def test_string_join(ldf, ignore_nulls, delimiter): ["abc", "def"], ], ) -def test_string_zfill(fill, input_strings): +def test_string_zfill(engine: pl.GPUEngine, fill, input_strings): ldf = pl.LazyFrame({"a": input_strings}) q = ldf.select(pl.col("a").str.zfill(fill)) @@ -566,7 +568,7 @@ def test_string_zfill(fill, input_strings): cudf_except=cudf_except, ) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -580,10 +582,10 @@ def test_string_zfill(fill, input_strings): else pytest.param(999, marks=pytest.mark.xfail(reason="fixed in Polars 1.30")), ], ) -def test_string_zfill_pl_129(fill): +def test_string_zfill_pl_129(engine: pl.GPUEngine, fill): ldf = pl.LazyFrame({"a": ["-1", "+2"]}) q = ldf.select(pl.col("a").str.zfill(fill)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -602,7 +604,7 @@ def test_string_zfill_pl_129(fill): pytest.param(None, marks=pytest.mark.xfail(reason="None dtype")), ], ) -def test_string_zfill_column(fill): +def test_string_zfill_column(engine: pl.GPUEngine, fill): ldf = pl.DataFrame( { "input_strings": ["1", "0", "123", "45", "", "0", "-1", "+2", "abc", "def"], @@ -620,7 +622,7 @@ def test_string_zfill_column(fill): cudf_except=cudf_except, ) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_string_zfill_forbidden_chars(): @@ -659,14 +661,14 @@ def test_string_zfill_forbidden_chars(): ), ], ) -def test_string_pad_start(width, char, using_rapidsmpf): - if using_rapidsmpf: +def test_string_pad_start(engine: pl.GPUEngine, width, char, using_streaming_engine): + if using_streaming_engine: pytest.skip( "Avoiding possible segfault with cuda 12.9 builds https://github.com/rapidsai/cudf/issues/21828" ) df = pl.LazyFrame({"a": ["abc", "defg", "hij"]}) q = df.select(pl.col("a").str.pad_start(width, char)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -695,26 +697,26 @@ def test_string_pad_start(width, char, using_rapidsmpf): ), ], ) -def test_string_pad_end(width, char): +def test_string_pad_end(engine: pl.GPUEngine, width, char): df = pl.LazyFrame({"a": ["abc", "defg", "hij"]}) q = df.select(pl.col("a").str.pad_end(width, char)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("ignore_nulls", [False, True]) @pytest.mark.parametrize("delimiter", ["", "-"]) -def test_string_join_non_string_data(ignore_nulls, delimiter): +def test_string_join_non_string_data(engine: pl.GPUEngine, ignore_nulls, delimiter): ldf = pl.LazyFrame({"a": [1, None, 3]}) q = ldf.select(pl.col("a").str.join(delimiter, ignore_nulls=ignore_nulls)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_reverse(ldf): +def test_string_reverse(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.reverse()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_to_titlecase(): +def test_string_to_titlecase(engine: pl.GPUEngine): df = pl.LazyFrame( { "quotes": [ @@ -727,43 +729,43 @@ def test_string_to_titlecase(): q = df.with_columns( quotes_title=pl.col("quotes").str.to_titlecase(), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("tail", [1, 2, 999, -1, 0, None]) -def test_string_tail(ldf, tail): +def test_string_tail(engine: pl.GPUEngine, ldf, tail): q = ldf.select(pl.col("a").str.tail(tail)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("head", [1, 2, 999, -1, 0, None]) -def test_string_head(ldf, head): +def test_string_head(engine: pl.GPUEngine, ldf, head): q = ldf.select(pl.col("a").str.head(head)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("ignore_nulls", [True, False]) @pytest.mark.parametrize("separator", ["*", ""]) -def test_concat_horizontal(ldf, ignore_nulls, separator): +def test_concat_horizontal(engine: pl.GPUEngine, ldf, ignore_nulls, separator): q = ldf.select( pl.concat_str(["a", "c"], separator=separator, ignore_nulls=ignore_nulls) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("ascii_case_insensitive", [True, False]) -def test_contains_any(ldf, ascii_case_insensitive): +def test_contains_any(engine: pl.GPUEngine, ldf, ascii_case_insensitive): q = ldf.select( pl.col("a").str.contains_any( ["a", "b", "c"], ascii_case_insensitive=ascii_case_insensitive ) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_count_matches(ldf): +def test_count_matches(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.count_matches("a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_count_matches_literal_unsupported(ldf): @@ -771,30 +773,30 @@ def test_count_matches_literal_unsupported(ldf): assert_ir_translation_raises(q, NotImplementedError) -def test_strip_prefix(ldf): +def test_strip_prefix(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.strip_prefix("A")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_strip_suffix(ldf): +def test_strip_suffix(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.strip_suffix("e")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_strip_prefix_suffix_dupes(): +def test_strip_prefix_suffix_dupes(engine: pl.GPUEngine): ldf = pl.LazyFrame({"a": ["a", "aa", "ab", "bb", "b"]}) q = ldf.select(pl.col("a").str.strip_prefix("a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) q = ldf.select(pl.col("a").str.strip_suffix("a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) q = ldf.select(pl.col("a").str.strip_prefix("b")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) q = ldf.select(pl.col("a").str.strip_suffix("b")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.fixture @@ -804,33 +806,33 @@ def ldf_jsonlike(): ) -def test_json_decode(ldf_jsonlike): +def test_json_decode(engine: pl.GPUEngine, ldf_jsonlike): q = ldf_jsonlike.select(pl.col("a").str.json_decode(pl.Struct({"a": pl.String()}))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) if POLARS_VERSION_LT_133: - q = ldf_jsonlike.select(pl.col("a").str.json_decode(None)) + q = ldf_jsonlike.select(pl.col("a").str.json_decode(None)) # type: ignore[arg-type] assert_ir_translation_raises(q, NotImplementedError) @pytest.mark.parametrize("dtype", [pl.Int64(), pl.Float64()]) -def test_json_decode_numeric_types(dtype): +def test_json_decode_numeric_types(engine: pl.GPUEngine, dtype): ldf = pl.LazyFrame({"a": ['{"a": 1}', None, '{"a": 2}']}) q = ldf.select(pl.col("a").str.json_decode(pl.Struct({"a": dtype}))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_json_decode_nested(): +def test_json_decode_nested(engine: pl.GPUEngine): ldf = pl.LazyFrame({"a": ['{"a": {"b": 1}}', None]}) q = ldf.select( pl.col("a").str.json_decode(pl.Struct({"a": pl.Struct({"b": pl.Int64()})})) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_json_path_match(ldf_jsonlike): +def test_json_path_match(engine: pl.GPUEngine, ldf_jsonlike): q = ldf_jsonlike.select(pl.col("a").str.json_path_match("$.a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.fixture @@ -850,9 +852,9 @@ def test_find_literal_false_strict_false_unsupported(ldf_find): @pytest.mark.parametrize("literal", [True, False]) @pytest.mark.parametrize("pattern", ["a|e", "a"]) -def test_find_literal(ldf_find, literal, pattern): +def test_find_literal(engine: pl.GPUEngine, ldf_find, literal, pattern): q = ldf_find.select(pl.col("a").str.find(pattern, literal=literal)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_find_literal_false_column_unsupported(ldf_find): @@ -866,9 +868,9 @@ def ldf_extract(): @pytest.mark.parametrize("group_index", [1, 2]) -def test_extract(ldf_extract, group_index): +def test_extract(engine: pl.GPUEngine, ldf_extract, group_index): q = ldf_extract.select(pl.col("a").str.extract(r"(\S+) (\d+) (.+)", group_index)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_extract_group_index_0_unsupported(ldf_extract): @@ -876,38 +878,38 @@ def test_extract_group_index_0_unsupported(ldf_extract): assert_ir_translation_raises(q, NotImplementedError) -def test_extract_groups(ldf_extract): +def test_extract_groups(engine: pl.GPUEngine, ldf_extract): q = ldf_extract.select(pl.col("a").str.extract_groups(r"(\S+) (\d+) (.+)")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_len_bytes(ldf): +def test_len_bytes(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.len_bytes()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_len_chars(ldf): +def test_len_chars(engine: pl.GPUEngine, ldf): q = ldf.select(pl.col("a").str.len_chars()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_string_concat_empty_frame(): +def test_string_concat_empty_frame(engine: pl.GPUEngine): lf = pl.LazyFrame({"a": pl.Series([], dtype=pl.String)}) q = lf.select(pl.lit(", ") + pl.col("a")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_single_column_concat_str(): +def test_single_column_concat_str(engine: pl.GPUEngine): lf = pl.LazyFrame({"c0": ["a", "b"]}) q = lf.select(pl.concat_str(pl.col("c0"))) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_concat_str_with_boolean(): +def test_concat_str_with_boolean(engine: pl.GPUEngine): lf = pl.LazyFrame({"c0": [True, False, None]}) q = lf.with_columns(pl.concat_str([pl.col("c0"), pl.lit("bool")])) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.skipif( diff --git a/python/cudf_polars/tests/expressions/test_struct.py b/python/cudf_polars/tests/expressions/test_struct.py index 2e4686b9079b..fd682201d7fe 100644 --- a/python/cudf_polars/tests/expressions/test_struct.py +++ b/python/cudf_polars/tests/expressions/test_struct.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -20,7 +20,7 @@ def ldf(): ) -def test_field_getitem(request, ldf): +def test_field_getitem(engine: pl.GPUEngine, request, ldf): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -28,11 +28,11 @@ def test_field_getitem(request, ldf): ) ) q = ldf.select(pl.col("a").struct[0]) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("fields", [("b",), ("b", "d"), ("^b.*|f.*$",)]) -def test_field(request, ldf, fields): +def test_field(engine: pl.GPUEngine, request, ldf, fields): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -40,10 +40,10 @@ def test_field(request, ldf, fields): ) ) q = ldf.select(pl.col("a").struct.field(*fields)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_unnest(request, ldf): +def test_unnest(engine: pl.GPUEngine, request, ldf): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -51,10 +51,10 @@ def test_unnest(request, ldf): ) ) q = ldf.select(pl.col("a").struct.unnest()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_json_encode(request, ldf): +def test_json_encode(engine: pl.GPUEngine, request, ldf): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -62,14 +62,14 @@ def test_json_encode(request, ldf): ) ) q = ldf.select(pl.col("a").struct.json_encode()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) ldf_newlines = pl.LazyFrame({"a": [{"b": "c\nd", "d": "\r\nz"}]}) q = ldf_newlines.select(pl.col("a").struct.json_encode()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_rename_fields(request, ldf): +def test_rename_fields(engine: pl.GPUEngine, request, ldf): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -77,7 +77,7 @@ def test_rename_fields(request, ldf): ) ) q = ldf.select(pl.col("a").struct.rename_fields(["1", "2", "3"]).struct.unnest()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_with_fields(ldf): @@ -92,7 +92,7 @@ def test_with_fields(ldf): [pl.col("a").name.prefix_fields, pl.col("a").name.suffix_fields], ids=lambda x: x.__name__, ) -def test_prefix_suffix_fields(request, ldf, expr): +def test_prefix_suffix_fields(engine: pl.GPUEngine, request, ldf, expr): request.applymarker( pytest.mark.xfail( condition=POLARS_VERSION_LT_131, @@ -100,7 +100,7 @@ def test_prefix_suffix_fields(request, ldf, expr): ) ) q = ldf.select(expr("foo").struct.unnest()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_map_field_names(ldf): @@ -110,16 +110,16 @@ def test_map_field_names(ldf): @pytest.mark.parametrize("name", [None, "my_count"]) @pytest.mark.parametrize("normalize", [True, False]) -def test_value_counts(ldf, name, normalize): +def test_value_counts(engine: pl.GPUEngine, ldf, name, normalize): # sort=True since order is non-deterministic q = ldf.select(pl.col("a").value_counts(sort=True, name=name, normalize=normalize)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_value_counts_normalize_div_by_zero(): +def test_value_counts_normalize_div_by_zero(engine: pl.GPUEngine): ldf = pl.LazyFrame({"a": []}, schema={"a": pl.Int64()}) q = ldf.select(pl.col("a").value_counts(normalize=True)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_groupby_value_counts_notimplemented(): @@ -132,18 +132,18 @@ def test_groupby_value_counts_notimplemented(): assert_ir_translation_raises(q, NotImplementedError) -def test_struct(ldf): +def test_struct(engine: pl.GPUEngine, ldf): q = ldf.select(pl.struct(pl.all())) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_nested_struct(): +def test_nested_struct(engine: pl.GPUEngine): ldf = pl.LazyFrame({"a": [{"x": {"i": 0, "j": 0}, "y": {"i": 0, "k": 1}}]}) q = ldf.select(pl.struct(pl.all())) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_value_counts_with_nulls(ldf): +def test_value_counts_with_nulls(engine: pl.GPUEngine, ldf): ldf_with_nulls = ldf.select(c=pl.Series(["x", None, "y", "x", None, "x"])) q = ldf_with_nulls.select(pl.col("c").value_counts(sort=True)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/expressions/test_top_and_bottom_k.py b/python/cudf_polars/tests/expressions/test_top_and_bottom_k.py index fdc7c7df6f69..5eb8f26c993e 100644 --- a/python/cudf_polars/tests/expressions/test_top_and_bottom_k.py +++ b/python/cudf_polars/tests/expressions/test_top_and_bottom_k.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -25,13 +25,13 @@ def df(): @pytest.mark.parametrize("col", ["test", "bool_val", "str_value", "col_with_nulls"]) @pytest.mark.parametrize("k", [0, 1, 2, 3, 4]) -def test_top_k(df, col, k): +def test_top_k(engine: pl.GPUEngine, df, col, k): q = df.select(pl.col(col).top_k(k)) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize("col", ["test", "bool_val", "str_value", "col_with_nulls"]) @pytest.mark.parametrize("k", [0, 1, 2, 3, 4]) -def test_bottom_k(df, col, k): +def test_bottom_k(engine: pl.GPUEngine, df, col, k): q = df.select(pl.col(col).bottom_k(k)) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) diff --git a/python/cudf_polars/tests/expressions/test_unique.py b/python/cudf_polars/tests/expressions/test_unique.py index 6b5886012e3c..94c4ecf8b2b7 100644 --- a/python/cudf_polars/tests/expressions/test_unique.py +++ b/python/cudf_polars/tests/expressions/test_unique.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("maintain_order", [False, True], ids=["unstable", "stable"]) @pytest.mark.parametrize("pre_sorted", [False, True], ids=["unsorted", "sorted"]) -def test_unique(maintain_order, pre_sorted): +def test_unique(engine: pl.GPUEngine, maintain_order, pre_sorted): ldf = pl.DataFrame( { "b": [1.5, 2.5, None, 1.5, 3, float("nan"), 3], @@ -21,11 +21,11 @@ def test_unique(maintain_order, pre_sorted): ldf = ldf.sort("b") query = ldf.select(pl.col("b").unique(maintain_order=maintain_order)) - assert_gpu_result_equal(query, check_row_order=maintain_order) + assert_gpu_result_equal(query, engine=engine, check_row_order=maintain_order) -def test_unique_on_sorted_expression(): +def test_unique_on_sorted_expression(engine: pl.GPUEngine): # Sort expr produces a column with is_sorted=YES, covering the sorted fast-path ldf = pl.DataFrame({"b": [3.0, 1.0, 2.0, 1.0, 3.0]}).lazy() query = ldf.select(pl.col("b").sort().unique()) - assert_gpu_result_equal(query, check_row_order=False) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) diff --git a/python/cudf_polars/tests/expressions/test_when_then.py b/python/cudf_polars/tests/expressions/test_when_then.py index 951df47eb571..b1300855cec2 100644 --- a/python/cudf_polars/tests/expressions/test_when_then.py +++ b/python/cudf_polars/tests/expressions/test_when_then.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 @@ -20,7 +20,7 @@ ) @pytest.mark.parametrize("then", [pl.lit(10), pl.col("a")]) @pytest.mark.parametrize("otherwise", [pl.lit(-2), pl.col("b")]) -def test_when_then(when, then, otherwise): +def test_when_then(engine: pl.GPUEngine, when, then, otherwise): df = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -30,4 +30,4 @@ def test_when_then(when, then, otherwise): ) q = df.select(pl.when(when).then(then).otherwise(otherwise)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/test_cache.py b/python/cudf_polars/tests/test_cache.py index c0a8c971cc2f..0a13ecb82f37 100644 --- a/python/cudf_polars/tests/test_cache.py +++ b/python/cudf_polars/tests/test_cache.py @@ -14,7 +14,7 @@ from cudf_polars.utils.versions import POLARS_VERSION_LT_1323 -def test_cache(request): +def test_cache(engine: pl.GPUEngine, request): request.applymarker( pytest.mark.xfail( condition=not POLARS_VERSION_LT_1323, @@ -30,7 +30,7 @@ def test_cache(request): df2 = pl.LazyFrame({"a": [7, 8], "b": [12, 13]}) q = pl.concat([df1, df2, df1, df2, df1]) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) t = Translator(q._ldf.visit(), pl.GPUEngine()) qir = t.translate_ir() diff --git a/python/cudf_polars/tests/test_dataframescan.py b/python/cudf_polars/tests/test_dataframescan.py index 2986c58b5fef..ba1acb1d0335 100644 --- a/python/cudf_polars/tests/test_dataframescan.py +++ b/python/cudf_polars/tests/test_dataframescan.py @@ -31,7 +31,7 @@ ], ) @pytest.mark.parametrize("predicate_pushdown", [False, True]) -def test_scan_drop_nulls(subset, predicate_pushdown): +def test_scan_drop_nulls(engine: pl.GPUEngine, subset, predicate_pushdown): df = pl.LazyFrame( { "a": [1, 2, 3, 4], @@ -46,13 +46,14 @@ def test_scan_drop_nulls(subset, predicate_pushdown): assert_gpu_result_equal( q, + engine=engine, collect_kwargs={ "optimizations": pl.QueryOptFlags(predicate_pushdown=predicate_pushdown) }, ) -def test_can_convert_lists(): +def test_can_convert_lists(engine: pl.GPUEngine): df = pl.LazyFrame( { "a": pl.Series([[1, 2], [3]], dtype=pl.List(pl.Int8())), @@ -68,10 +69,10 @@ def test_can_convert_lists(): } ) - assert_gpu_result_equal(df) + assert_gpu_result_equal(df, engine=engine) -def test_dataframescan_with_decimals(): +def test_dataframescan_with_decimals(engine: pl.GPUEngine): q = pl.LazyFrame( { "foo": [1, 2], @@ -79,23 +80,25 @@ def test_dataframescan_with_decimals(): }, schema={"foo": pl.Int64, "bar": pl.Decimal(precision=15, scale=2)}, ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.skipif( POLARS_VERSION_LT_138, reason="height parameter added in Polars 1.38", ) -def test_dataframescan_zero_width_with_rows(request, using_rapidsmpf): +def test_dataframescan_zero_width_with_rows( + engine: pl.GPUEngine, request, using_streaming_engine +): request.applymarker( pytest.mark.xfail( - using_rapidsmpf, + using_streaming_engine, reason="https://github.com/rapidsai/cudf/issues/21644", ) ) df = pl.LazyFrame(height=5) q = df.select(pl.len()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_struct_literal_not_supported(): diff --git a/python/cudf_polars/tests/test_distinct.py b/python/cudf_polars/tests/test_distinct.py index d42c4a96f5a0..514fadc23e66 100644 --- a/python/cudf_polars/tests/test_distinct.py +++ b/python/cudf_polars/tests/test_distinct.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -13,7 +13,7 @@ @pytest.mark.parametrize("keep", ["any", "none", "first", "last"]) @pytest.mark.parametrize("maintain_order", [False, True], ids=["unstable", "stable"]) @pytest.mark.parametrize("pre_sorted", [False, True], ids=["unsorted", "sorted"]) -def test_distinct(subset, keep, maintain_order, pre_sorted): +def test_distinct(engine: pl.GPUEngine, subset, keep, maintain_order, pre_sorted): ldf = pl.DataFrame( { "a": [1, 2, 1, 3, 5, None, None], @@ -27,4 +27,4 @@ def test_distinct(subset, keep, maintain_order, pre_sorted): ldf = ldf.sort(*keys, descending=descending) query = ldf.unique(subset=subset, keep=keep, maintain_order=maintain_order) - assert_gpu_result_equal(query, check_row_order=maintain_order) + assert_gpu_result_equal(query, engine=engine, check_row_order=maintain_order) diff --git a/python/cudf_polars/tests/test_drop_nulls.py b/python/cudf_polars/tests/test_drop_nulls.py index 0fe9b963eddb..b2f1a2098536 100644 --- a/python/cudf_polars/tests/test_drop_nulls.py +++ b/python/cudf_polars/tests/test_drop_nulls.py @@ -32,9 +32,9 @@ def null_data(request): ).lazy() -def test_drop_null(null_data): +def test_drop_null(engine: pl.GPUEngine, null_data): q = null_data.select(pl.col("a").drop_nulls()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -42,41 +42,41 @@ 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(engine: pl.GPUEngine, null_data, value): q = null_data.select(pl.col("a").fill_null(value)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_fill_null_with_string(): +def test_fill_null_with_string(engine: pl.GPUEngine): q = pl.LazyFrame({"a": [None, "a"]}).select(pl.col("a").fill_null("b")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( "strategy", ["forward", "backward", "min", "max", "mean", "zero", "one"] ) -def test_fill_null_with_strategy(null_data, strategy): +def test_fill_null_with_strategy(engine: pl.GPUEngine, null_data, strategy): q = null_data.select(pl.col("a").fill_null(strategy=strategy)) if POLARS_VERSION_LT_132: assert_ir_translation_raises(q, NotImplementedError) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("strategy", ["zero", "one"]) -def test_fill_null_with_strategy_bool(strategy): +def test_fill_null_with_strategy_bool(engine: pl.GPUEngine, strategy): q = pl.LazyFrame({"a": [True, None, False]}).select( pl.col("a").fill_null(strategy=strategy) ) if POLARS_VERSION_LT_132: assert_ir_translation_raises(q, NotImplementedError) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("strategy", ["forward", "backward"]) @pytest.mark.parametrize("limit", [0, 1, 2]) -def test_fill_null_with_limit(null_data, strategy, limit): +def test_fill_null_with_limit(engine: pl.GPUEngine, null_data, strategy, limit): q = null_data.select(pl.col("a").fill_null(strategy=strategy, limit=limit)) if limit != 0: assert_ir_translation_raises(q, NotImplementedError) @@ -84,4 +84,4 @@ def test_fill_null_with_limit(null_data, strategy, limit): if POLARS_VERSION_LT_132: assert_ir_translation_raises(q, NotImplementedError) else: - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/test_filter.py b/python/cudf_polars/tests/test_filter.py index 64b143729d68..20dc656470f6 100644 --- a/python/cudf_polars/tests/test_filter.py +++ b/python/cudf_polars/tests/test_filter.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("expr", [pl.col("c"), pl.col("b") < 1, pl.lit(value=True)]) @pytest.mark.parametrize("predicate_pushdown", [False, True]) -def test_filter(expr, predicate_pushdown): +def test_filter(engine: pl.GPUEngine, expr, predicate_pushdown): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -23,6 +23,7 @@ def test_filter(expr, predicate_pushdown): query = ldf.filter(expr) assert_gpu_result_equal( query, + engine=engine, collect_kwargs={ "optimizations": pl.QueryOptFlags(predicate_pushdown=predicate_pushdown) }, diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index 5686e824fcc2..af7fa5772c0a 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -126,20 +126,27 @@ def maintain_order(request): return request.param -def test_groupby(df: pl.LazyFrame, maintain_order, keys, exprs): +def test_groupby(engine: pl.GPUEngine, df: pl.LazyFrame, maintain_order, keys, exprs): q = df.group_by(*keys, maintain_order=maintain_order).agg(*exprs) if not maintain_order: sort_keys = list(q.collect_schema().keys())[: len(keys)] q = q.sort(*sort_keys) - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) -def test_groupby_sorted_keys(df: pl.LazyFrame, keys, exprs, using_rapidsmpf, request): +def test_groupby_sorted_keys( + engine: pl.GPUEngine, + df: pl.LazyFrame, + keys, + exprs, + using_streaming_engine, + request, +): request.applymarker( pytest.mark.xfail( - using_rapidsmpf, + using_streaming_engine, strict=False, reason="https://github.com/rapidsai/cudf/issues/21642 - no deterministic sort for keys", ) @@ -159,18 +166,18 @@ def test_groupby_sorted_keys(df: pl.LazyFrame, keys, exprs, using_rapidsmpf, req # 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, check_exact=False) + 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, check_exact=False) + assert_gpu_result_equal(qsorted, engine=engine, check_exact=False) else: - assert_gpu_result_equal(q, check_exact=False) + assert_gpu_result_equal(q, engine=engine, check_exact=False) -def test_groupby_len(df, keys): +def test_groupby_len(engine: pl.GPUEngine, df, keys): q = df.group_by(*keys).agg(pl.len()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -193,7 +200,7 @@ def test_groupby_unsupported(df: pl.LazyFrame, expr: pl.Expr) -> None: assert_ir_translation_raises(q, NotImplementedError) -def test_groupby_null_keys(maintain_order): +def test_groupby_null_keys(engine: pl.GPUEngine, maintain_order): df = pl.LazyFrame( { "key": pl.Series([1, float("nan"), 2, None, 2, None], dtype=pl.Float64()), @@ -205,11 +212,11 @@ def test_groupby_null_keys(maintain_order): if not maintain_order: q = q.sort("key") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.xfail(reason="https://github.com/pola-rs/polars/issues/17513") -def test_groupby_minmax_with_nan(): +def test_groupby_minmax_with_nan(engine: pl.GPUEngine): df = pl.LazyFrame( {"key": [1, 2, 2, 2], "value": [float("nan"), 1, -1, float("nan")]} ) @@ -218,7 +225,7 @@ def test_groupby_minmax_with_nan(): pl.col("value").max().alias("max"), pl.col("value").min().alias("min") ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("op", [pl.Expr.nan_max, pl.Expr.nan_min]) @@ -252,11 +259,11 @@ def test_groupby_nan_minmax_raises(op): [pl.lit(2).alias("value"), pl.col("float") * 2], ], ) -def test_groupby_literal_in_agg(df, key, expr): +def test_groupby_literal_in_agg(engine: pl.GPUEngine, df, key, expr): # check_row_order=False doesn't work for list aggregations # so just sort by the group key q = df.group_by(key).agg(expr).sort(key, maintain_order=True) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -294,7 +301,16 @@ def test_groupby_nested_list_struct_raises(dtype): @pytest.mark.parametrize("nrows", [30, 300, 300_000]) @pytest.mark.parametrize("nkeys", [1, 2, 4]) -def test_groupby_maintain_order_random(nrows, nkeys, with_nulls): +def test_groupby_maintain_order_random( + engine: pl.GPUEngine, + blocksize_mode, + nrows, + nkeys, + with_nulls, + using_streaming_engine, +): + if nrows > 30 and (blocksize_mode == "small" or using_streaming_engine): + pytest.skip("streaming executor too slow for large n_rows") key_names = [f"key{key}" for key in range(nkeys)] rng = random.Random(2) key_values = [rng.choices(range(100), k=nrows) for _ in key_names] @@ -311,34 +327,33 @@ def test_groupby_maintain_order_random(nrows, nkeys, with_nulls): ) ) q = df.lazy().group_by(key_names, maintain_order=True).agg(pl.col("value").sum()) - # The streaming executor is too slow for large n_rows with blocksize_mode="small" - assert_gpu_result_equal(q, blocksize_mode="default" if nrows > 30 else None) + assert_gpu_result_equal(q, engine=engine) -def test_groupby_len_with_nulls(): +def test_groupby_len_with_nulls(engine: pl.GPUEngine): df = pl.DataFrame({"a": [1, 1, 1, 2], "b": [1, None, 2, 3]}) q = df.lazy().group_by("a").agg(pl.col("b").len()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize("column", ["int", "string", "uint16_with_null"]) -def test_groupby_nunique(df: pl.LazyFrame, column): +def test_groupby_nunique(engine: pl.GPUEngine, df: pl.LazyFrame, column): q = df.group_by("key1").agg(pl.col(column).n_unique()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize("column", ["int", "string", "uint16_with_null"]) -def test_groupby_nunique_drop_nulls(df: pl.LazyFrame, column): +def test_groupby_nunique_drop_nulls(engine: pl.GPUEngine, df: pl.LazyFrame, column): q = df.group_by("key1").agg(pl.col(column).drop_nulls().n_unique()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) -def test_groupby_null_count(df: pl.LazyFrame): +def test_groupby_null_count(engine: pl.GPUEngine, df: pl.LazyFrame): q = df.group_by("key1").agg(pl.col("uint16_with_null").null_count()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -365,15 +380,15 @@ def test_groupby_unsupported_non_pointwise_boolean_function(df: pl.LazyFrame, ex assert_ir_translation_raises(q, NotImplementedError) -def test_groupby_mean_type_promotion(df: pl.LazyFrame) -> None: +def test_groupby_mean_type_promotion(engine: pl.GPUEngine, df: pl.LazyFrame) -> None: df = df.with_columns(pl.col("float").cast(pl.Float32)) q = df.group_by("key1").agg(pl.col("float").mean()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) -def test_groupby_sum_all_null_group_returns_null(): +def test_groupby_sum_all_null_group_returns_null(engine: pl.GPUEngine): df = pl.LazyFrame( { "key": ["a", "a", "b", "b", "c"], @@ -382,7 +397,7 @@ def test_groupby_sum_all_null_group_returns_null(): ) q = df.group_by("key").agg(out=pl.col("null_groups").sum()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -396,7 +411,7 @@ def test_groupby_sum_all_null_group_returns_null(): ids=["sum", "mean", "median", "quantile-0.5"], ) def test_groupby_aggs_keep_unsupported_as_null( - request, df: pl.LazyFrame, agg_expr + engine: pl.GPUEngine, request, df: pl.LazyFrame, agg_expr ) -> None: request.applymarker( pytest.mark.xfail( @@ -406,7 +421,7 @@ def test_groupby_aggs_keep_unsupported_as_null( ) lf = df.filter(pl.col("datetime") == date(2004, 12, 1)) q = lf.group_by("datetime").agg(agg_expr) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -437,9 +452,11 @@ def test_groupby_aggs_keep_unsupported_as_null( "post_manually_compute_mean", ], ) -def test_groupby_ternary_supported(df: pl.LazyFrame, expr: pl.Expr) -> None: +def test_groupby_ternary_supported( + engine: pl.GPUEngine, df: pl.LazyFrame, expr: pl.Expr +) -> None: q = df.group_by("key1").agg(expr) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -464,23 +481,23 @@ def test_groupby_rank_raises(df: pl.LazyFrame) -> None: assert_ir_translation_raises(q, NotImplementedError) -def test_groupby_sum_decimal_null_group() -> None: +def test_groupby_sum_decimal_null_group(engine: pl.GPUEngine) -> None: df = pl.LazyFrame( {"key1": [1, 1, 2, 3], "foo": [None, None, Decimal("1.00"), Decimal("2.00")]}, schema={"key1": pl.Int32, "foo": pl.Decimal(9, 2)}, ) q = df.group_by("key1").agg(pl.col("foo").sum()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.xfail( raises=AssertionError, reason="https://github.com/rapidsai/cudf/issues/19610", ) -def test_groupby_literal_agg(): +def test_groupby_literal_agg(engine: pl.GPUEngine): df = pl.LazyFrame({"c0": [True, False]}) q = df.group_by("c0").agg(pl.lit(1).is_not_null()) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) def test_groupby_empty_keys_raises(): diff --git a/python/cudf_polars/tests/test_hconcat.py b/python/cudf_polars/tests/test_hconcat.py index 2da57c02a679..761b673a682a 100644 --- a/python/cudf_polars/tests/test_hconcat.py +++ b/python/cudf_polars/tests/test_hconcat.py @@ -9,7 +9,7 @@ from cudf_polars.testing.asserts import assert_gpu_result_equal -def test_hconcat(): +def test_hconcat(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -18,16 +18,16 @@ def test_hconcat(): ).lazy() ldf2 = ldf.select((pl.col("a") + pl.col("b")).alias("c")) query = pl.concat([ldf, ldf2], how="horizontal") - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_hconcat_different_heights(): +def test_hconcat_different_heights(engine: pl.GPUEngine): left = pl.LazyFrame({"a": [1, 2, 3, 4]}) right = pl.LazyFrame({"b": [[1], [2]], "c": ["a", "bcde"]}) q = pl.concat([left, right], how="horizontal") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_hconcat_should_broadcast(): diff --git a/python/cudf_polars/tests/test_hstack.py b/python/cudf_polars/tests/test_hstack.py index b8c97f4607f1..7aafafbcfbb6 100644 --- a/python/cudf_polars/tests/test_hstack.py +++ b/python/cudf_polars/tests/test_hstack.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -7,7 +7,7 @@ from cudf_polars.testing.asserts import assert_gpu_result_equal -def test_hstack(): +def test_hstack(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -16,10 +16,10 @@ def test_hstack(): ).lazy() query = ldf.with_columns(pl.col("a") + pl.col("b")) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_hstack_with_cse(): +def test_hstack_with_cse(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -29,4 +29,4 @@ def test_hstack_with_cse(): expr = pl.col("a") + pl.col("b") query = ldf.with_columns(expr.alias("c"), expr.alias("d") * 2) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) diff --git a/python/cudf_polars/tests/test_join.py b/python/cudf_polars/tests/test_join.py index 404fca50087f..a513552512e2 100644 --- a/python/cudf_polars/tests/test_join.py +++ b/python/cudf_polars/tests/test_join.py @@ -14,10 +14,7 @@ from cudf_polars.containers import DataType 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, - get_default_engine, -) +from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.versions import POLARS_VERSION_LT_132 @@ -61,10 +58,10 @@ def right(): @pytest.mark.parametrize( "maintain_order", ["left", "left_right", "right_left", "right"] ) -def test_join_maintain_order(left, right, maintain_order): +def test_join_maintain_order(engine: pl.GPUEngine, left, right, maintain_order): q = left.join(right, on=pl.col("a"), how="inner", maintain_order=maintain_order) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -77,11 +74,18 @@ def test_join_maintain_order(left, right, maintain_order): ], ) def test_non_coalesce_join( - left, right, how, nulls_equal, join_expr, using_rapidsmpf, request + engine: pl.GPUEngine, + left, + right, + how, + nulls_equal, + join_expr, + using_streaming_engine, + request, ): request.applymarker( pytest.mark.xfail( - using_rapidsmpf, + using_streaming_engine, strict=False, reason="Non deterministic sort/join on nulls", ) @@ -89,7 +93,7 @@ def test_non_coalesce_join( query = left.join( right, on=join_expr, how=how, nulls_equal=nulls_equal, coalesce=False ) - assert_gpu_result_equal(query, check_row_order=False) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -99,14 +103,14 @@ def test_non_coalesce_join( ["c", "a"], ], ) -def test_coalesce_join(left, right, how, nulls_equal, join_expr): +def test_coalesce_join(engine: pl.GPUEngine, left, right, how, nulls_equal, join_expr): query = left.join( right, on=join_expr, how=how, nulls_equal=nulls_equal, coalesce=True ) - assert_gpu_result_equal(query, check_row_order=False) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) -def test_left_join_with_slice(left, right, nulls_equal, zlice): +def test_left_join_with_slice(engine: pl.GPUEngine, left, right, nulls_equal, zlice): q = left.join(right, on="a", how="left", nulls_equal=nulls_equal, coalesce=True) if zlice is not None: @@ -116,25 +120,23 @@ def test_left_join_with_slice(left, right, nulls_equal, zlice): # the things that invariant to the ordering. q = q.slice(*zlice) - engine = get_default_engine() - # Check the number of rows - assert_gpu_result_equal(q.select(pl.len())) + assert_gpu_result_equal(q.select(pl.len()), engine=engine) # Check that the schema matches result = q.collect(engine=engine) assert result.schema == q.collect_schema() else: - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) -def test_cross_join(left, right, zlice): +def test_cross_join(engine: pl.GPUEngine, left, right, zlice): q = left.join(right, how="cross") if zlice is not None: q = q.slice(*zlice) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -144,9 +146,9 @@ def test_cross_join(left, right, zlice): (pl.lit(2, dtype=pl.Int64), pl.col("a")), ], ) -def test_join_literal_key(left, right, left_on, right_on): +def test_join_literal_key(engine: pl.GPUEngine, left, right, left_on, right_on): q = left.join(right, left_on=left_on, right_on=right_on, how="inner") - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) @pytest.mark.parametrize( @@ -163,20 +165,23 @@ def test_join_literal_key(left, right, left_on, right_on): ], ) @pytest.mark.parametrize("zlice", [None, (0, 5)]) -def test_join_where(left, right, conditions, zlice): +@pytest.mark.skip_on_streaming_engine( + "ConditionalJoin not supported for multiple partitions" +) +def test_join_where(engine: pl.GPUEngine, left, right, conditions, zlice): q = left.join_where(right, *conditions) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) if zlice is not None: q_len = q.slice(*zlice).select(pl.len()) # Can't compare result, since row order is not guaranteed and # therefore we only check the length - assert_gpu_result_equal(q_len) + assert_gpu_result_equal(q_len, engine=engine) -def test_cross_join_empty_right_table(request): +def test_cross_join_empty_right_table(engine: pl.GPUEngine, request): request.applymarker( pytest.mark.xfail(condition=POLARS_VERSION_LT_132, reason="nested loop join") ) @@ -187,20 +192,24 @@ def test_cross_join_empty_right_table(request): (pl.col("a") == pl.col("a")) & (pl.col("b") < pl.col("b")) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("maintain_order", ["left_right", "right_left"]) @pytest.mark.parametrize("how", ["inner", "full"]) -def test_join_maintain_order_inner_full(left, right, how, maintain_order, nulls_equal): +def test_join_maintain_order_inner_full( + engine: pl.GPUEngine, left, right, how, maintain_order, nulls_equal +): q = left.join( right, on="a", how=how, nulls_equal=nulls_equal, maintain_order=maintain_order ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("maintain_order", ["left", "left_right"]) -def test_join_maintain_order_left(left, right, maintain_order, nulls_equal): +def test_join_maintain_order_left( + engine: pl.GPUEngine, left, right, maintain_order, nulls_equal +): q = left.join( right, on="a", @@ -208,11 +217,13 @@ def test_join_maintain_order_left(left, right, maintain_order, nulls_equal): nulls_equal=nulls_equal, maintain_order=maintain_order, ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("maintain_order", ["right", "right_left"]) -def test_join_maintain_order_right(left, right, maintain_order, nulls_equal): +def test_join_maintain_order_right( + engine: pl.GPUEngine, left, right, maintain_order, nulls_equal +): q = left.join( right, on="a", @@ -220,15 +231,17 @@ def test_join_maintain_order_right(left, right, maintain_order, nulls_equal): nulls_equal=nulls_equal, maintain_order=maintain_order, ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("maintain_order", ["left_right", "right_left"]) @pytest.mark.parametrize("join_expr", [pl.col("a"), ["c", "a"]]) @pytest.mark.parametrize("how", ["inner", "full"]) -def test_join_maintain_order_multiple_keys(left, right, how, join_expr, maintain_order): +def test_join_maintain_order_multiple_keys( + engine: pl.GPUEngine, left, right, how, join_expr, maintain_order +): q = left.join(right, on=join_expr, how=how, maintain_order=maintain_order) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -242,9 +255,11 @@ def test_join_maintain_order_multiple_keys(left, right, how, join_expr, maintain ("right_left", "full"), ], ) -def test_join_maintain_order_with_coalesce(left, right, maintain_order, how): +def test_join_maintain_order_with_coalesce( + engine: pl.GPUEngine, left, right, maintain_order, how +): q = left.join(right, on="a", how=how, coalesce=True, maintain_order=maintain_order) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -256,12 +271,15 @@ def test_join_maintain_order_with_coalesce(left, right, maintain_order, how): ("right_left", "full", (2, 3)), ], ) -def test_join_maintain_order_with_slice(left, right, maintain_order, how, zlice): +def test_join_maintain_order_with_slice( + engine: pl.GPUEngine, left, right, maintain_order, how, zlice +): # Need to disable slice pushdown to make the test deterministic. We want to materialize # the full join result and then slice q = left.join(right, on="a", how=how, maintain_order=maintain_order).slice(*zlice) assert_gpu_result_equal( q, + engine=engine, polars_collect_kwargs={"optimizations": pl.QueryOptFlags(slice_pushdown=False)}, ) @@ -290,7 +308,12 @@ def test_join_maintain_order_with_slice(left, right, maintain_order, how, zlice) (pl.Decimal(15, 2), pl.Float64), ], ) -def test_cross_join_filter_with_decimals(request, expr, left_dtype, right_dtype): +@pytest.mark.skip_on_streaming_engine( + "ConditionalJoin not supported for multiple partitions" +) +def test_cross_join_filter_with_decimals( + engine: pl.GPUEngine, request, expr, left_dtype, right_dtype +): request.applymarker( pytest.mark.xfail( POLARS_VERSION_LT_132 @@ -324,7 +347,7 @@ def test_cross_join_filter_with_decimals(request, expr, left_dtype, right_dtype) q = left.join(right, how="cross").filter(expr) - assert_gpu_result_equal(q, check_row_order=False) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) def test_conditional_join_predicate_pickle(): diff --git a/python/cudf_polars/tests/test_mapfunction.py b/python/cudf_polars/tests/test_mapfunction.py index 1590bb00ef1d..4b57b06de4df 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -24,7 +24,7 @@ def test_explode_multiple_raises(): @pytest.mark.parametrize("column", ["a", "b"]) -def test_explode_single(column): +def test_explode_single(engine: pl.GPUEngine, column): df = pl.LazyFrame( { "a": [[1, 2], [3, 4], None], @@ -34,7 +34,7 @@ def test_explode_single(column): ) q = df.explode(column) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("mapping", [{"b": "a"}, {"a": "c", "b": "c"}]) @@ -54,18 +54,18 @@ def test_rename_duplicate_raises(mapping): @pytest.mark.parametrize( "mapping", [{}, {"b": "c"}, {"b": "a", "a": "b"}, {"a": "c", "b": "d"}] ) -def test_rename_columns(mapping): +def test_rename_columns(engine: pl.GPUEngine, mapping): df = pl.LazyFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) q = df.rename(mapping) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("index", [None, ["a"], ["d", "a"]]) @pytest.mark.parametrize("variable_name", [None, "names"]) @pytest.mark.parametrize("value_name", [None, "unpivoted"]) -def test_unpivot(index, variable_name, value_name): +def test_unpivot(engine: pl.GPUEngine, index, variable_name, value_name): df = pl.LazyFrame( { "a": ["x", "y", "z"], @@ -78,10 +78,10 @@ def test_unpivot(index, variable_name, value_name): ["c", "b"], index=index, variable_name=variable_name, value_name=value_name ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_unpivot_defaults(): +def test_unpivot_defaults(engine: pl.GPUEngine): df = pl.LazyFrame( { "a": pl.Series([11, 12, 13], dtype=pl.UInt16), @@ -91,10 +91,10 @@ def test_unpivot_defaults(): } ) q = df.unpivot(index="d") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_with_row_index_defaults(): +def test_with_row_index_defaults(engine: pl.GPUEngine): lf = pl.LazyFrame( { "a": [1, 3, 5], @@ -102,7 +102,7 @@ def test_with_row_index_defaults(): } ) q = lf.with_row_index() - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_unique_hash(): @@ -115,7 +115,7 @@ def test_unique_hash(): assert hash(ir_a) != hash(ir_b) -def test_set_sorted_then_inner_join(request): +def test_set_sorted_then_inner_join(engine: pl.GPUEngine, request): request.applymarker( pytest.mark.xfail( condition=not POLARS_VERSION_LT_135, @@ -127,7 +127,7 @@ def test_set_sorted_then_inner_join(request): q = df.set_sorted("a").join( pl.LazyFrame({"a": [2, 4], "b": [20, 40]}), on="a", how="inner" ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_explode_single_legacy_options(): diff --git a/python/cudf_polars/tests/test_merge_sorted.py b/python/cudf_polars/tests/test_merge_sorted.py index ab654d8445eb..1c3448b37234 100644 --- a/python/cudf_polars/tests/test_merge_sorted.py +++ b/python/cudf_polars/tests/test_merge_sorted.py @@ -10,10 +10,12 @@ @pytest.mark.parametrize("descending", [True, False]) -def test_merge_sorted_without_nulls(descending, request, using_rapidsmpf): +def test_merge_sorted_without_nulls( + engine: pl.GPUEngine, descending, request, using_streaming_engine +): request.applymarker( pytest.mark.xfail( - not using_rapidsmpf and descending, + not using_streaming_engine and descending, reason="https://github.com/pola-rs/polars/issues/21511", ) ) @@ -28,7 +30,7 @@ def test_merge_sorted_without_nulls(descending, request, using_rapidsmpf): } ).sort("age", descending=descending) q = df0.merge_sorted(df1, key="age") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -41,7 +43,7 @@ def test_merge_sorted_without_nulls(descending, request, using_rapidsmpf): False, ], ) -def test_merge_sorted_with_nulls(descending): +def test_merge_sorted_with_nulls(engine: pl.GPUEngine, descending): df0 = pl.LazyFrame( { "name": ["steve", "elise", "bob", "john"], @@ -57,4 +59,4 @@ def test_merge_sorted_with_nulls(descending): } ).sort("age", descending=descending) q = df0.merge_sorted(df1, key="age") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/test_parquet_filters.py b/python/cudf_polars/tests/test_parquet_filters.py index 98b6514d3761..c789a953d1e9 100644 --- a/python/cudf_polars/tests/test_parquet_filters.py +++ b/python/cudf_polars/tests/test_parquet_filters.py @@ -60,8 +60,8 @@ def test_scan_by_hand(expr, selection, pq_file, chunked): ) -def test_parquet_filter_boolean_column(tmp_path): +def test_parquet_filter_boolean_column(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"x": [1, 2, 3], "y": [True, False, True]}) df.write_parquet(tmp_path / "df.parquet") q = pl.scan_parquet(tmp_path / "df.parquet").filter(pl.col("y")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/test_rolling.py b/python/cudf_polars/tests/test_rolling.py index a9a0617ff719..40de49360f73 100644 --- a/python/cudf_polars/tests/test_rolling.py +++ b/python/cudf_polars/tests/test_rolling.py @@ -51,7 +51,7 @@ def df(presort): @pytest.mark.parametrize("closed", ["left", "right", "both", "none"]) @pytest.mark.parametrize("period", ["1w4d", "48h", "180s"]) -def test_datetime_rolling(df, closed, period): +def test_datetime_rolling(engine: pl.GPUEngine, df, closed, period): q = df.rolling("dt", period=period, closed=closed).agg( sum_a=pl.sum("values"), min_a=pl.min("values"), @@ -59,11 +59,11 @@ def test_datetime_rolling(df, closed, period): count=pl.len(), ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("dtype", [pl.Int32, pl.UInt32, pl.Int64, pl.UInt64]) -def test_rolling_integral_orderby(dtype): +def test_rolling_integral_orderby(engine: pl.GPUEngine, dtype): df = pl.LazyFrame( { "orderby": pl.Series([1, 4, 8, 10, 12, 13, 14, 22], dtype=dtype), @@ -72,7 +72,7 @@ def test_rolling_integral_orderby(dtype): ) q = df.rolling("orderby", period="4i", closed="both").agg(pl.col("values").sum()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_rolling_collect_list_raises(): @@ -89,7 +89,7 @@ def test_rolling_collect_list_raises(): @pytest.mark.parametrize("with_slice", [False, True]) -def test_rolling_empty_aggs(with_slice): +def test_rolling_empty_aggs(engine: pl.GPUEngine, with_slice): df = pl.LazyFrame( { "orderby": [1, 4, 8, 10, 12, 13, 14, 22], @@ -100,7 +100,7 @@ def test_rolling_empty_aggs(with_slice): if with_slice: q = q.slice(2) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_calendrical_period_unsupported(df): @@ -109,7 +109,7 @@ def test_calendrical_period_unsupported(df): assert_ir_translation_raises(q, NotImplementedError) -def test_unsorted_raises(): +def test_unsorted_raises(engine_raise_on_fail: pl.GPUEngine): df = pl.LazyFrame({"orderby": [1, 2, 4, 2], "values": [1, 2, 3, 4]}) q = df.rolling("orderby", period="2i").agg(sum=pl.sum("values")) with pytest.raises(pl.exceptions.InvalidOperationError): @@ -117,10 +117,10 @@ def test_unsorted_raises(): with pytest.raises( RuntimeError, match=r".*rolling is not sorted, please sort first" ): - q.collect(engine=pl.GPUEngine(raise_on_fail=True)) + q.collect(engine=engine_raise_on_fail) -def test_grouped_rolling(): +def test_grouped_rolling(engine: pl.GPUEngine): df = pl.LazyFrame( { "keys": [1, None, 2, 1, 2, None], @@ -130,10 +130,10 @@ def test_grouped_rolling(): ) q = df.rolling("orderby", period="5i", group_by="keys").agg(pl.col("values").sum()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_grouped_rolling_unsorted_raises(): +def test_grouped_rolling_unsorted_raises(engine_raise_on_fail: pl.GPUEngine): df = pl.LazyFrame( { "keys": [1, None, 2, 1, 2, None], @@ -146,10 +146,10 @@ def test_grouped_rolling_unsorted_raises(): with pytest.raises(pl.exceptions.ComputeError): q.collect(engine="in-memory") with pytest.raises(RuntimeError, match="Input for grouped rolling is not sorted"): - q.collect(engine=pl.GPUEngine(raise_on_fail=True)) + q.collect(engine=engine_raise_on_fail) -def test_orderby_nulls_raises_computeerror(): +def test_orderby_nulls_raises_computeerror(engine_raise_on_fail: pl.GPUEngine): df = pl.LazyFrame({"orderby": [1, 2, 4, None], "values": [1, 2, 3, 4]}) q = df.rolling("orderby", period="2i").agg(sum=pl.sum("values")) with pytest.raises(pl.exceptions.InvalidOperationError): @@ -157,7 +157,7 @@ def test_orderby_nulls_raises_computeerror(): with pytest.raises( RuntimeError, match=r"Index column.*in rolling may not contain nulls" ): - q.collect(engine=pl.GPUEngine(raise_on_fail=True)) + q.collect(engine=engine_raise_on_fail) def test_rolling_nested_raises(request): @@ -196,7 +196,7 @@ def test_unsupported_agg(): assert_ir_translation_raises(q, NotImplementedError) -def test_rolling_sum_all_null_window_returns_null(): +def test_rolling_sum_all_null_window_returns_null(engine: pl.GPUEngine): df = pl.LazyFrame( { "orderby": [1, 2, 3, 4, 5, 6], @@ -207,17 +207,17 @@ def test_rolling_sum_all_null_window_returns_null(): out=pl.col("null_windows").sum() ) # Expected: [0, 0, 5, 5, 5, 1] - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_rolling_null_count(df): +def test_rolling_null_count(engine: pl.GPUEngine, df): lf = df.with_columns( null=pl.when(pl.col("values") % 2 == 0).then(None).otherwise(pl.col("values")) ) q = lf.rolling("dt", period="48h", closed="both").agg( nc=pl.col("null").null_count() ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -244,9 +244,9 @@ def test_rolling_null_count(df): "post_manually_compute_mean", ], ) -def test_rolling_ternary_supported(df, expr): +def test_rolling_ternary_supported(engine: pl.GPUEngine, df, expr): q = df.rolling("dt", period="48h", closed="both").agg(expr.alias("out")) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 66f74b9bce12..0f58893bc4f9 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -184,7 +184,7 @@ def test_scan_row_index_projected_out(tmp_path): assert_gpu_result_equal(q, engine=NO_CHUNK_ENGINE) -def test_scan_csv_column_renames_projection_schema(tmp_path): +def test_scan_csv_column_renames_projection_schema(engine: pl.GPUEngine, tmp_path): with (tmp_path / "test.csv").open("w") as f: f.write("""foo,bar,baz\n1,2\n3,4,5""") @@ -198,7 +198,7 @@ def test_scan_csv_column_renames_projection_schema(tmp_path): }, ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -221,7 +221,7 @@ def test_scan_csv_column_renames_projection_schema(tmp_path): (4, 2), ], ) -def test_scan_csv_multi(tmp_path, filename, glob, nrows_skiprows): +def test_scan_csv_multi(engine: pl.GPUEngine, tmp_path, filename, glob, nrows_skiprows): n_rows, skiprows = nrows_skiprows with (tmp_path / "test1.csv").open("w") as f: f.write("""foo,bar,baz\n1,2,3\n3,4,5""") @@ -235,7 +235,7 @@ def test_scan_csv_multi(tmp_path, filename, glob, nrows_skiprows): source = tmp_path / filename q = pl.scan_csv(source, glob=glob, n_rows=n_rows, skip_rows=skiprows) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_scan_csv_multi_differing_colnames(tmp_path): @@ -277,35 +277,35 @@ def test_scan_csv_comment_str_not_implemented(tmp_path): assert_ir_translation_raises(q, NotImplementedError) -def test_scan_csv_comment_char(tmp_path): +def test_scan_csv_comment_char(engine: pl.GPUEngine, tmp_path): with (tmp_path / "test.csv").open("w") as f: f.write("""foo,bar,baz\n# 1,2,3\n3,4,5""") q = pl.scan_csv(tmp_path / "test.csv", comment_prefix="#") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize("nulls", [None, "3", ["3", "5"]]) -def test_scan_csv_null_values(tmp_path, nulls): +def test_scan_csv_null_values(engine: pl.GPUEngine, tmp_path, nulls): with (tmp_path / "test.csv").open("w") as f: f.write("""foo,bar,baz\n1,2,3\n3,4,5\n5,,2""") q = pl.scan_csv(tmp_path / "test.csv", null_values=nulls) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_scan_csv_decimal_comma(tmp_path): +def test_scan_csv_decimal_comma(engine: pl.GPUEngine, tmp_path): with (tmp_path / "test.csv").open("w") as f: f.write("""foo|bar|baz\n1,23|2,34|3,56\n1""") q = pl.scan_csv(tmp_path / "test.csv", separator="|", decimal_comma=True) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_scan_csv_skip_initial_empty_rows(tmp_path): +def test_scan_csv_skip_initial_empty_rows(engine: pl.GPUEngine, tmp_path): with (tmp_path / "test.csv").open("w") as f: f.write("""\n\n\n\nfoo|bar|baz\n1|2|3\n1""") @@ -315,16 +315,16 @@ def test_scan_csv_skip_initial_empty_rows(tmp_path): q = pl.scan_csv(tmp_path / "test.csv", separator="|", skip_rows=1) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_scan_csv_slice_end_none(tmp_path): +def test_scan_csv_slice_end_none(engine: pl.GPUEngine, tmp_path): with (tmp_path / "test.csv").open("w") as f: f.write("""c0\ntrue\nfalse""") q = pl.scan_csv(tmp_path / "test.csv").slice(10, None) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -335,10 +335,10 @@ def test_scan_csv_slice_end_none(tmp_path): {"a": pl.UInt64}, ], ) -def test_scan_ndjson_schema(df, tmp_path, schema): +def test_scan_ndjson_schema(engine: pl.GPUEngine, df, tmp_path, schema): make_partitioned_source(df, tmp_path / "file", "ndjson") q = pl.scan_ndjson(tmp_path / "file", schema=schema) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_scan_ndjson_unsupported(df, tmp_path): @@ -439,13 +439,13 @@ def test_scan_parquet_chunked( ) -def test_select_arbitrary_order_with_row_index_column(tmp_path): +def test_select_arbitrary_order_with_row_index_column(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"a": [1, 2, 3]}) df.write_parquet(tmp_path / "df.parquet") q = pl.scan_parquet(tmp_path / "df.parquet", row_index_name="foo").select( [pl.col("a"), pl.col("foo")] ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -456,7 +456,14 @@ def test_select_arbitrary_order_with_row_index_column(tmp_path): ], ) def test_scan_csv_with_and_without_header( - df, tmp_path, has_header, new_columns, row_index, columns, zlice + engine: pl.GPUEngine, + df, + tmp_path, + has_header, + new_columns, + row_index, + columns, + zlice, ): path = tmp_path / "test.csv" make_partitioned_source( @@ -478,7 +485,7 @@ def test_scan_csv_with_and_without_header( if columns is not None: q = q.select(columns) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_scan_csv_without_header_and_new_column_names_raises(df, tmp_path): @@ -488,13 +495,13 @@ def test_scan_csv_without_header_and_new_column_names_raises(df, tmp_path): assert_ir_translation_raises(q, NotImplementedError) -def test_scan_with_row_index(tmp_path: Path) -> None: +def test_scan_with_row_index(engine: pl.GPUEngine, tmp_path: Path) -> None: df = pl.DataFrame({"a": [1, 2, 3, 4]}) df.write_csv(tmp_path / "test-0.csv") df.write_csv(tmp_path / "test-1.csv") q = pl.scan_csv(tmp_path / "test-*.csv", row_index_name="index", row_index_offset=0) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) def test_scan_from_file_uri(tmp_path: Path) -> None: @@ -573,6 +580,7 @@ def get_handler(req: Request) -> Response: def test_scan_ndjson_remote( + engine: pl.GPUEngine, request: pytest.FixtureRequest, tmp_path: Path, df: pl.DataFrame, @@ -617,10 +625,12 @@ def get_handler(_: Request) -> Response: ) q = pl.scan_ndjson(httpserver.url_for(server_path)) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_scan_parquet_with_decimal_literal_in_predicate(df, tmp_path): +def test_scan_parquet_with_decimal_literal_in_predicate( + engine: pl.GPUEngine, df, tmp_path +): make_partitioned_source(df, tmp_path / "file", "parquet") q = pl.scan_parquet(tmp_path / "file").filter( @@ -628,20 +638,20 @@ def test_scan_parquet_with_decimal_literal_in_predicate(df, tmp_path): & (pl.lit(Decimal("2.00")).cast(pl.Decimal(15, 2)) < pl.col("d")) ) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_scan_csv_blank_line(tmp_path): +def test_scan_csv_blank_line(engine: pl.GPUEngine, tmp_path): data = """c0 polars""" fle = tmp_path / "test.csv" fle.write_text(data) q = pl.scan_csv(fle) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_hits_scan_row_index_duplicate(request, tmp_path): +def test_hits_scan_row_index_duplicate(engine: pl.GPUEngine, request, tmp_path): request.applymarker( pytest.mark.xfail( condition=not POLARS_VERSION_LT_138, @@ -656,7 +666,7 @@ def test_hits_scan_row_index_duplicate(request, tmp_path): if POLARS_VERSION_LT_135: # Did not raise before - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) else: assert_ir_translation_raises(q, NotImplementedError) @@ -687,29 +697,33 @@ def test_scan_compressed_file_raises(tmp_path, compression, file_type): assert_ir_translation_raises(q, NotImplementedError) -def test_scan_tiny_file_not_compressed(tmp_path): +def test_scan_tiny_file_not_compressed(engine: pl.GPUEngine, tmp_path): # code coverage for the case where we try to # detect compression but the file is too small # to have a valid signature. path = tmp_path / "tiny.csv" path.write_bytes(b"a\n") q = pl.scan_csv(path, has_header=False, new_columns=["a"]) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.skipif( POLARS_VERSION_LT_138, reason="height parameter added in Polars 1.38", ) -@pytest.mark.parametrize("engine", [None, NO_CHUNK_ENGINE]) -def test_scan_parquet_zero_width_with_limit(tmp_path, engine, request, using_rapidsmpf): +@pytest.mark.parametrize("custom_engine", [None, NO_CHUNK_ENGINE]) +def test_scan_parquet_zero_width_with_limit( + engine: pl.GPUEngine, tmp_path, custom_engine, request, using_streaming_engine +): request.applymarker( pytest.mark.xfail( - using_rapidsmpf and engine is None, + using_streaming_engine and custom_engine is None, reason="https://github.com/rapidsai/cudf/issues/21644", ) ) path = tmp_path / "zero_width.parquet" pl.LazyFrame(height=20).sink_parquet(path) q = pl.scan_parquet(path).head(5) - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal( + q, engine=custom_engine if custom_engine is not None else engine + ) diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index ac9dbb90d0c0..16ae1cf4384c 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -15,7 +15,7 @@ from cudf_polars.utils.versions import POLARS_VERSION_LT_134 -def test_select(): +def test_select(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -27,15 +27,15 @@ def test_select(): pl.col("a") + pl.col("b"), (pl.col("a") * 2 + pl.col("b")).alias("d") ) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_select_decimal(): +def test_select_decimal(engine: pl.GPUEngine): ldf = pl.LazyFrame( {"a": pl.Series(values=[decimal.Decimal("1.0"), None], dtype=pl.Decimal(3, 1))} ) query = ldf.select(pl.col("a")) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) def test_select_decimal_precision_none_result_max_precision(): @@ -55,7 +55,7 @@ def test_select_decimal_precision_none_result_max_precision(): assert gpu_result.schema["a"].precision == 38 -def test_select_reduce(): +def test_select_reduce(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -68,26 +68,26 @@ def test_select_reduce(): (pl.col("a") * 2 + pl.col("b")).alias("d").mean(), ) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize("expr", [pl.col("a").first(), pl.col("a").last()]) -def test_select_first_last_empty(expr): +def test_select_first_last_empty(engine: pl.GPUEngine, expr): ldf = pl.LazyFrame({"a": []}, schema={"a": pl.Int64}) query = ldf.select(expr) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_select_with_cse_no_agg(): +def test_select_with_cse_no_agg(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) expr = pl.col("a") + pl.col("a") query = df.select(expr, (expr * 2).alias("b"), ((expr * 2) + 10).alias("c")) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_select_with_cse_with_agg(): +def test_select_with_cse_with_agg(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) expr = pl.col("a") + pl.col("a") asum = pl.col("a").sum() + pl.col("a").sum() @@ -96,13 +96,13 @@ def test_select_with_cse_with_agg(): expr, (expr * 2).alias("b"), asum.alias("c"), (asum + 10).alias("d") ) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_select_native_datetime(): +def test_select_native_datetime(engine: pl.GPUEngine): df = pl.LazyFrame({"c0": [1]}) query = df.select(pl.datetime(1969, 12, 7, 20, 47, 14)) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) @pytest.mark.parametrize("fmt", ["ndjson", "csv"]) @@ -122,13 +122,13 @@ def test_select_fast_count_unsupported_formats(tmp_path, fmt): assert_ir_translation_raises(q, NotImplementedError) -def test_select_fast_count_parquet(tmp_path): +def test_select_fast_count_parquet(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"a": [1, 2, 3]}) file = tmp_path / "data.parquet" df.write_parquet(file) q = pl.scan_parquet(file).select(pl.len()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) @pytest.mark.parametrize( @@ -139,10 +139,12 @@ def test_select_fast_count_parquet(tmp_path): (-1,), ], ) -def test_select_fast_count_parquet_skip_rows(request, tmp_path, zlice): +def test_select_fast_count_parquet_skip_rows( + engine: pl.GPUEngine, request, tmp_path, zlice +): df = pl.DataFrame({"a": [1, 2, 3]}) file = tmp_path / "data.parquet" df.write_parquet(file) q = pl.scan_parquet(file).slice(1, 5).select(pl.len()) - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/cudf_polars/tests/test_sink.py b/python/cudf_polars/tests/test_sink.py index 022d1a9962cb..e10e248ab1c6 100644 --- a/python/cudf_polars/tests/test_sink.py +++ b/python/cudf_polars/tests/test_sink.py @@ -7,7 +7,6 @@ import polars as pl from cudf_polars.testing.asserts import ( - DEFAULT_BLOCKSIZE_MODE, assert_sink_ir_translation_raises, assert_sink_result_equal, ) @@ -28,13 +27,23 @@ def df(): @pytest.mark.parametrize("null_value", [None, "NA"]) @pytest.mark.parametrize("line_terminator", ["\n", "\n\n"]) @pytest.mark.parametrize("separator", [",", "|"]) -def test_sink_csv(df, tmp_path, include_header, null_value, line_terminator, separator): - if line_terminator == "\n\n" and DEFAULT_BLOCKSIZE_MODE == "small": +def test_sink_csv( + engine: pl.GPUEngine, + blocksize_mode, + df, + tmp_path, + include_header, + null_value, + line_terminator, + separator, +): + if line_terminator == "\n\n" and blocksize_mode == "small": # We end up with an extra row per partition. pytest.skip("Multi-line terminator not supported with small blocksize") assert_sink_result_equal( df, tmp_path / "out.csv", + engine=engine, write_kwargs={ "include_header": include_header, "null_value": null_value, @@ -69,10 +78,11 @@ def test_sink_csv_unsupported_kwargs(df, tmp_path, kwarg, value): ) -def test_sink_ndjson(df, tmp_path): +def test_sink_ndjson(engine: pl.GPUEngine, df, tmp_path): assert_sink_result_equal( df, tmp_path / "out.ndjson", + engine=engine, ) diff --git a/python/cudf_polars/tests/test_slice.py b/python/cudf_polars/tests/test_slice.py index 7208f4fbe747..1bcfe1ac3596 100644 --- a/python/cudf_polars/tests/test_slice.py +++ b/python/cudf_polars/tests/test_slice.py @@ -18,7 +18,7 @@ [0, 2, 12, 11], ) @pytest.mark.parametrize("slice_pushdown", [False, True]) -def test_slice(offset, length, slice_pushdown): +def test_slice(engine: pl.GPUEngine, offset, length, slice_pushdown): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -34,6 +34,7 @@ def test_slice(offset, length, slice_pushdown): ) assert_gpu_result_equal( query, + engine=engine, collect_kwargs={ "optimizations": pl.QueryOptFlags(slice_pushdown=slice_pushdown) }, diff --git a/python/cudf_polars/tests/test_sort.py b/python/cudf_polars/tests/test_sort.py index cfa8e5ff9b95..f93828c7c6f7 100644 --- a/python/cudf_polars/tests/test_sort.py +++ b/python/cudf_polars/tests/test_sort.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -20,7 +20,7 @@ ) @pytest.mark.parametrize("nulls_last", [False, True]) @pytest.mark.parametrize("maintain_order", [False, True], ids=["unstable", "stable"]) -def test_sort(sort_keys, nulls_last, maintain_order): +def test_sort(engine: pl.GPUEngine, sort_keys, nulls_last, maintain_order): ldf = pl.DataFrame( { "a": [1, 2, 1, 3, 5, None, None], @@ -36,4 +36,4 @@ def test_sort(sort_keys, nulls_last, maintain_order): nulls_last=nulls_last, maintain_order=maintain_order, ) - assert_gpu_result_equal(query, check_row_order=maintain_order) + assert_gpu_result_equal(query, engine=engine, check_row_order=maintain_order) diff --git a/python/cudf_polars/tests/test_union.py b/python/cudf_polars/tests/test_union.py index a1d1e9c5b989..e855f2856521 100644 --- a/python/cudf_polars/tests/test_union.py +++ b/python/cudf_polars/tests/test_union.py @@ -9,7 +9,7 @@ ) -def test_union(): +def test_union(engine: pl.GPUEngine): ldf = pl.DataFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -18,10 +18,10 @@ def test_union(): ).lazy() ldf2 = ldf.select((pl.col("a") + pl.col("b")).alias("c"), pl.col("a")) query = pl.concat([ldf, ldf2], how="diagonal") - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) -def test_concat_vertical(): +def test_concat_vertical(engine: pl.GPUEngine): ldf = pl.LazyFrame( { "a": [1, 2, 3, 4, 5, 6, 7], @@ -31,10 +31,10 @@ def test_concat_vertical(): ldf2 = ldf.select(pl.col("a"), pl.col("b") * 2 + pl.col("a")) q = pl.concat([ldf, ldf2], how="vertical") - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) -def test_concat_diagonal_empty(): +def test_concat_diagonal_empty(engine: pl.GPUEngine): df1 = pl.LazyFrame() df2 = pl.LazyFrame({"a": [1, 2]}) @@ -42,5 +42,6 @@ def test_concat_diagonal_empty(): assert_gpu_result_equal( q, + engine=engine, collect_kwargs={"optimizations": pl.QueryOptFlags()}, ) diff --git a/python/cudf_polars/tests/test_window_functions.py b/python/cudf_polars/tests/test_window_functions.py index e7f9181f833c..cf7f9181652c 100644 --- a/python/cudf_polars/tests/test_window_functions.py +++ b/python/cudf_polars/tests/test_window_functions.py @@ -77,7 +77,7 @@ def unsupported_agg_expr(request): return request.param -def test_over(df: pl.LazyFrame, partition_by, agg_expr): +def test_over(engine: pl.GPUEngine, df: pl.LazyFrame, partition_by, agg_expr): """Test window functions over partitions.""" window_expr = agg_expr.over(partition_by) @@ -91,8 +91,8 @@ def test_over(df: pl.LazyFrame, partition_by, agg_expr): # GPU: 1.333333333333334 # Classic floating-point gotcha: looks the same, but the test fails assert_gpu_result_equal( - q, check_exact=False, rtol=1e-15, atol=1e-15 - ) if "var" in str(agg_expr) else assert_gpu_result_equal(q) + q, engine=engine, check_exact=False, rtol=1e-15, atol=1e-15 + ) if "var" in str(agg_expr) else assert_gpu_result_equal(q, engine=engine) def test_over_with_sort(df: pl.LazyFrame): @@ -102,7 +102,9 @@ def test_over_with_sort(df: pl.LazyFrame): @pytest.mark.parametrize("mapping_strategy", ["group_to_rows", "explode", "join"]) -def test_over_mapping_strategy(df: pl.LazyFrame, mapping_strategy: str): +def test_over_mapping_strategy( + engine: pl.GPUEngine, df: pl.LazyFrame, mapping_strategy: str +): """Test window functions with different mapping strategies.""" # ignore is for polars' WindowMappingStrategy, which isn't publicly exported. # https://github.com/pola-rs/polars/issues/17420 @@ -119,13 +121,15 @@ def test_over_mapping_strategy(df: pl.LazyFrame, mapping_strategy: str): ] ) if not POLARS_VERSION_LT_132 and mapping_strategy == "group_to_rows": - assert_gpu_result_equal(q) + assert_gpu_result_equal(q, engine=engine) else: assert_ir_translation_raises(q, NotImplementedError) @pytest.mark.parametrize("period", ["2d", "3d"]) -def test_rolling(request, df: pl.LazyFrame, agg_expr, period: str): +def test_rolling( + engine: pl.GPUEngine, request, df: pl.LazyFrame, agg_expr, period: str +): """Test rolling window functions over time series.""" if not POLARS_VERSION_LT_136: request.applymarker( @@ -137,7 +141,7 @@ def test_rolling(request, df: pl.LazyFrame, agg_expr, period: str): query = df.with_columns(window_expr) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) def test_rolling_unsupported(df: pl.LazyFrame, unsupported_agg_expr): @@ -152,7 +156,7 @@ def test_rolling_unsupported(df: pl.LazyFrame, unsupported_agg_expr): @pytest.mark.parametrize("closed", ["left", "right", "both", "none"]) -def test_rolling_closed(request, df: pl.LazyFrame, closed: str): +def test_rolling_closed(engine: pl.GPUEngine, request, df: pl.LazyFrame, closed: str): """Test rolling window functions with different closed parameters.""" if not POLARS_VERSION_LT_136: request.applymarker( @@ -171,4 +175,4 @@ def test_rolling_closed(request, df: pl.LazyFrame, closed: str): ) ] ) - assert_gpu_result_equal(query) + assert_gpu_result_equal(query, engine=engine) diff --git a/python/cudf_polars/tests/testing/test_asserts.py b/python/cudf_polars/tests/testing/test_asserts.py index 16250c339984..ee5279181f1e 100644 --- a/python/cudf_polars/tests/testing/test_asserts.py +++ b/python/cudf_polars/tests/testing/test_asserts.py @@ -23,7 +23,7 @@ ) -def test_translation_assert_raises(): +def test_translation_assert_raises(engine: pl.GPUEngine): df = pl.LazyFrame( { "time": pl.datetime_range( @@ -37,7 +37,7 @@ def test_translation_assert_raises(): ) # This should succeed - assert_gpu_result_equal(df) + assert_gpu_result_equal(df, engine=engine) with pytest.raises(AssertionError): # This should fail, because we can translate this query.