fix(polars): use engine="streaming" instead of streaming=True

Author: cpcloud

ibis/backends/polars/__init__.py

@@ -441,8 +441,7 @@ def _to_dataframe(
         expr: ir.Expr,
         params: Mapping[ir.Expr, object] | None = None,
         limit: int | None = None,
-        streaming: bool = False,
-        engine: Literal["cpu", "gpu"] | pl.GPUEngine = "cpu",
+        engine: Literal["cpu", "gpu", "streaming"] | pl.GPUEngine = "cpu",
         **kwargs: Any,
     ) -> pl.DataFrame:
         self._run_pre_execute_hooks(expr)
@@ -452,7 +451,7 @@ def _to_dataframe(
             limit = ibis.options.sql.default_limit
         if limit is not None:
             lf = lf.limit(limit)
-        df = lf.collect(streaming=streaming, engine=engine)
+        df = lf.collect(engine=engine)
         # XXX: Polars sometimes returns data with the incorrect column names.
         # For now we catch this case and rename them here if needed.
         expected_cols = tuple(table_expr.columns)
@@ -467,17 +466,11 @@ def execute(
         *,
         params: Mapping[ir.Expr, object] | None = None,
         limit: int | None = None,
-        streaming: bool = False,
-        engine: Literal["cpu", "gpu"] | pl.GPUEngine = "cpu",
+        engine: Literal["cpu", "gpu", "streaming"] | pl.GPUEngine = "cpu",
         **kwargs: Any,
     ):
         df = self._to_dataframe(
-            expr,
-            params=params,
-            limit=limit,
-            streaming=streaming,
-            engine=engine,
-            **kwargs,
+            expr, params=params, limit=limit, engine=engine, **kwargs
         )
         if isinstance(expr, (ir.Table, ir.Scalar)):
             return expr.__pandas_result__(df.to_pandas())
@@ -500,17 +493,11 @@ def to_polars(
         *,
         params: Mapping[ir.Expr, object] | None = None,
         limit: int | None = None,
-        streaming: bool = False,
-        engine: Literal["cpu", "gpu"] | pl.GPUEngine = "cpu",
+        engine: Literal["cpu", "gpu", "streaming"] | pl.GPUEngine = "cpu",
         **kwargs: Any,
     ):
         df = self._to_dataframe(
-            expr,
-            params=params,
-            limit=limit,
-            streaming=streaming,
-            engine=engine,
-            **kwargs,
+            expr, params=params, limit=limit, engine=engine, **kwargs
         )
         return expr.__polars_result__(df)
 
@@ -519,19 +506,13 @@ def _to_pyarrow_table(
         expr: ir.Expr,
         params: Mapping[ir.Expr, object] | None = None,
         limit: int | None = None,
-        streaming: bool = False,
-        engine: Literal["cpu", "gpu"] | pl.GPUEngine = "cpu",
+        engine: Literal["cpu", "gpu", "streaming"] | pl.GPUEngine = "cpu",
         **kwargs: Any,
     ):
         from ibis.formats.pyarrow import PyArrowData
 
         df = self._to_dataframe(
-            expr,
-            params=params,
-            limit=limit,
-            streaming=streaming,
-            engine=engine,
-            **kwargs,
+            expr, params=params, limit=limit, engine=engine, **kwargs
         )
         return PyArrowData.convert_table(df.to_arrow(), expr.as_table().schema())
 

ibis/backends/polars/tests/test_client.py

@@ -66,13 +66,13 @@ def test_memtable_polars_types(con):
 def test_streaming(con, mocker, to_method):
     t = con.table("functional_alltypes")
     mocked_collect = mocker.patch("polars.LazyFrame.collect")
-    getattr(con, to_method)(t, streaming=True)
-    mocked_collect.assert_called_once_with(streaming=True, engine="cpu")
+    getattr(con, to_method)(t, engine="streaming")
+    mocked_collect.assert_called_once_with(engine="streaming")
 
 
 @pytest.mark.parametrize("to_method", ["to_pyarrow", "to_polars"])
 def test_engine(con, mocker, to_method):
     t = con.table("functional_alltypes")
     mocked_collect = mocker.patch("polars.LazyFrame.collect")
     getattr(con, to_method)(t, engine="gpu")
-    mocked_collect.assert_called_once_with(streaming=False, engine="gpu")
+    mocked_collect.assert_called_once_with(engine="gpu")