feat(rust,python): Add a config option to specify the default engine to use during lazyframe collect calls

crates/polars-core/src/config.rs

@@ -33,6 +33,13 @@ pub fn verbose() -> bool {
     std::env::var("POLARS_VERBOSE").as_deref().unwrap_or("") == "1"
 }
 
+pub fn use_gpu_engine() -> bool {
+    std::env::var("POLARS_DEFAULT_ENGINE")
+        .as_deref()
+        .unwrap_or("")
+        == "gpu"
+}
+
 pub fn get_file_prefetch_size() -> usize {
     std::env::var("POLARS_PREFETCH_SIZE")
         .map(|s| s.parse::<usize>().expect("integer"))

crates/polars-python/src/functions/utils.rs

@@ -1,7 +1,13 @@
 use polars::prelude::_set_check_length;
+use polars_core::config::use_gpu_engine as _use_gpu_engine;
 use pyo3::prelude::*;
 
 #[pyfunction]
 pub fn check_length(check: bool) {
     unsafe { _set_check_length(check) }
 }
+
+#[pyfunction]
+pub fn use_gpu_engine() -> PyResult<Option<bool>> {
+    Ok(Some(_use_gpu_engine()))
+}

py-polars/polars/config.py

@@ -71,6 +71,7 @@
     "POLARS_TABLE_WIDTH",
     "POLARS_VERBOSE",
     "POLARS_MAX_EXPR_DEPTH",
+    "POLARS_DEFAULT_ENGINE",
 }
 
 # vars that set the rust env directly should declare themselves here as the Config
@@ -140,6 +141,7 @@ class ConfigParameters(TypedDict, total=False):
     set_trim_decimal_zeros: bool | None
     set_verbose: bool | None
     set_expr_depth_warning: int
+    set_gpu_engine: bool | None
 
 
 class Config(contextlib.ContextDecorator):
@@ -1450,3 +1452,40 @@ def set_expr_depth_warning(cls, limit: int) -> type[Config]:
 
         os.environ["POLARS_MAX_EXPR_DEPTH"] = str(limit)
         return cls
+
+    @classmethod
+    def set_gpu_engine(cls, active: bool | None = None) -> type[Config]:
+        """
+        Set the default engine to use the GPU.
+
+        Parameters
+        ----------
+        engine : bool
+            Whether or not to default to GPU polars on .collect() calls
+
+        Examples
+        --------
+        >>> pl.Config.set_gpu_engine(True)  # doctest: +SKIP
+        >>> pl.Config.set_verbose(True)  # doctest: +SKIP
+        >>> lf = pl.DataFrame({"v": [1, 2, 3], "v2": [4, 5, 6]})  # doctest: +SKIP
+        >>> lf.max().collect()  # doctest: +SKIP
+        run PythonScanExec
+        run UdfExec
+        run StackExec
+        shape: (3, 2)
+        ┌─────┬─────┐
+        │ v   ┆ v2  │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 1   ┆ 4   │
+        │ 2   ┆ 5   │
+        │ 3   ┆ 6   │
+        └─────┴─────┘
+        """
+        if active is None:
+            os.environ.pop("POLARS_DEFAULT_ENGINE", None)
+        else:
+            engine = "gpu" if active else "cpu"
+            os.environ["POLARS_DEFAULT_ENGINE"] = engine
+        return cls

py-polars/polars/lazyframe/frame.py

@@ -81,6 +81,7 @@
 from polars.lazyframe.engine_config import GPUEngine
 from polars.lazyframe.group_by import LazyGroupBy
 from polars.lazyframe.in_process import InProcessQuery
+from polars.polars import use_gpu_engine
 from polars.schema import Schema
 from polars.selectors import by_dtype, expand_selector
 
@@ -2003,6 +2004,8 @@ def collect(
         if not (is_config_obj or engine in ("cpu", "gpu")):
             msg = f"Invalid engine argument {engine=}"
             raise ValueError(msg)
+        if use_gpu_engine():
+            is_gpu = True
         if (streaming or background or new_streaming) and is_gpu:
             issue_warning(
                 "GPU engine does not support streaming or background collection, "

py-polars/src/lib.rs

@@ -220,6 +220,8 @@ fn polars(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     // Functions: other
     m.add_wrapped(wrap_pyfunction!(functions::check_length))
         .unwrap();
+    m.add_wrapped(wrap_pyfunction!(functions::use_gpu_engine))
+        .unwrap();
 
     #[cfg(feature = "sql")]
     m.add_wrapped(wrap_pyfunction!(functions::sql_expr))

py-polars/tests/unit/test_config.py

@@ -965,6 +965,7 @@ def test_warn_unstable(recwarn: pytest.WarningsRecorder) -> None:
 @pytest.mark.parametrize(
     ("environment_variable", "config_setting", "value", "expected"),
     [
+        ("POLARS_DEFAULT_ENGINE", "set_gpu_engine", True, "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"),