Refactor cudf-polars test suite onto pytest fixtures

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

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

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 {}

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)