address review

pola-rs · Feb 5, 2025 · b4d9c5c · b4d9c5c
1 parent 6a7e1ea
commit b4d9c5c
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/py-polars/polars/config.py b/py-polars/polars/config.py
@@ -141,7 +141,7 @@ class ConfigParameters(TypedDict, total=False):
     set_trim_decimal_zeros: bool | None
     set_verbose: bool | None
     set_expr_depth_warning: int
-    set_default_engine: Literal["cpu", "gpu", "streaming"] | None
+    set_engine_affinity: Literal["cpu", "gpu", "streaming"] | None
 
 
 class Config(contextlib.ContextDecorator):
@@ -1454,22 +1454,43 @@ def set_expr_depth_warning(cls, limit: int) -> type[Config]:
         return cls
 
     @classmethod
-    def set_default_engine(
+    def set_engine_affinity(
         cls, engine: Literal["cpu", "gpu", "streaming"] | None
     ) -> type[Config]:
         """
         Set which engine to use by default.
 
         Parameters
         ----------
-        engine : Literal["cpu", "gpu"]
+        engine : Literal["cpu", "gpu", "streaming"]
             The default engine to use on all .collect() calls
 
         Examples
         --------
-        >>> pl.Config.set_default_engine("gpu")  # doctest: +SKIP
         >>> pl.Config.set_verbose(True)  # doctest: +SKIP
-        >>> lf = pl.DataFrame({"v": [1, 2, 3], "v2": [4, 5, 6]})  # doctest: +SKIP
+        >>> pl.Config.set_engine_affinity("streaming")  # doctest: +SKIP
+        >>> lf.max().collect()  # doctest: +SKIP
+        polars-stream: updating graph state
+        polars-stream: running in_memory_source in subgraph
+        polars-stream: running reduce in subgraph
+        async thread count: 4
+        polars-stream: done running graph phase
+        polars-stream: updating graph state
+        polars-stream: running in_memory_sink in subgraph
+        polars-stream: running select in subgraph
+        polars-stream: running reduce in subgraph
+        polars-stream: done running graph phase
+        polars-stream: updating graph state
+        shape: (1, 2)
+        ┌─────┬─────┐
+        │ v   ┆ v2  │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 3   ┆ 6   │
+        └─────┴─────┘
+        >>> pl.Config.set_engine_affinity("gpu")  # doctest: +SKIP
+        >>> lf = pl.LazyFrame({"v": [1, 2, 3], "v2": [4, 5, 6]})  # doctest: +SKIP
         >>> lf.max().collect()  # doctest: +SKIP
         run PythonScanExec
         run UdfExec

diff --git a/py-polars/tests/unit/test_config.py b/py-polars/tests/unit/test_config.py
@@ -937,7 +937,7 @@ def test_warn_unstable(recwarn: pytest.WarningsRecorder) -> None:
 @pytest.mark.parametrize(
     ("environment_variable", "config_setting", "value", "expected"),
     [
-        ("POLARS_DEFAULT_ENGINE", "set_default_engine", "gpu", "gpu"),
+        ("POLARS_ENGINE_AFFINITY", "set_engine_affinity", "gpu", "gpu"),
         ("POLARS_AUTO_STRUCTIFY", "set_auto_structify", True, "1"),
         ("POLARS_FMT_MAX_COLS", "set_tbl_cols", 12, "12"),
         ("POLARS_FMT_MAX_ROWS", "set_tbl_rows", 3, "3"),