diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py index 1067d3879..0c1929fa6 100644 --- a/narwhals/_duckdb/dataframe.py +++ b/narwhals/_duckdb/dataframe.py @@ -81,19 +81,21 @@ def __getitem__(self, item: str) -> DuckDBInterchangeSeries: ) @overload - def collect(self, return_type: Literal["pyarrow"] = "pyarrow") -> ArrowDataFrame: ... + def collect( + self, eager_backend: Literal["pyarrow"] = "pyarrow" + ) -> ArrowDataFrame: ... @overload - def collect(self, return_type: Literal["pandas"]) -> PandasLikeDataFrame: ... + def collect(self, eager_backend: Literal["pandas"]) -> PandasLikeDataFrame: ... @overload - def collect(self, return_type: Literal["polars"]) -> PolarsDataFrame: ... + def collect(self, eager_backend: Literal["polars"]) -> PolarsDataFrame: ... def collect( self, - return_type: Literal["pyarrow", "pandas", "polars"] = "pyarrow", + eager_backend: Literal["pyarrow", "pandas", "polars"] = "pyarrow", ) -> ArrowDataFrame | PandasLikeDataFrame | PolarsDataFrame: - if return_type == "pyarrow": + if eager_backend == "pyarrow": try: import pyarrow as pa # ignore-banned-import except ModuleNotFoundError as exc: # pragma: no cover @@ -110,7 +112,7 @@ def collect( version=self._version, ) - elif return_type == "pandas": + elif eager_backend == "pandas": import pandas as pd # ignore-banned-import from narwhals._pandas_like.dataframe import PandasLikeDataFrame @@ -123,7 +125,7 @@ def collect( version=self._version, ) - elif return_type == "polars": + elif eager_backend == "polars": import polars as pl # ignore-banned-import from narwhals._polars.dataframe import PolarsDataFrame @@ -137,8 +139,8 @@ def collect( else: msg = ( - "Only the following `return_type`'s are supported: pyarrow, pandas and " - f"polars. Found '{return_type}'." + "Only the following `eager_backend`'s are supported: pyarrow, pandas and " + f"polars. Found '{eager_backend}'." ) raise ValueError(msg) diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index 4a26da20b..fce645fee 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -3633,7 +3633,7 @@ def collect( duckdb_kwargs: Allows to specify in which eager backend to materialize a DuckDBPyRelation backed LazyFrame. It is possible to choose among `pyarrow`, `pandas` or `polars` by declaring - `duckdb_kwargs={"return_type": ""}`. + `duckdb_kwargs={"eager_backend": ""}`. Returns: DataFrame diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py index 1ce9be958..73d1aaa7e 100644 --- a/narwhals/stable/v1/__init__.py +++ b/narwhals/stable/v1/__init__.py @@ -436,7 +436,7 @@ def collect( duckdb_kwargs: Allows to specify in which eager backend to materialize a DuckDBPyRelation backed LazyFrame. It is possible to choose among `pyarrow`, `pandas` or `polars` by declaring - `duckdb_kwargs={"return_type": ""}`. + `duckdb_kwargs={"eager_backend": ""}`. Returns: DataFrame diff --git a/tests/frame/collect_test.py b/tests/frame/collect_test.py index 84bf8dd63..2700499ab 100644 --- a/tests/frame/collect_test.py +++ b/tests/frame/collect_test.py @@ -23,7 +23,7 @@ def test_collect_kwargs(constructor: Constructor) -> None: .collect( polars_kwargs={"no_optimization": True}, dask_kwargs={"optimize_graph": False}, - duckdb_kwargs={"return_type": "pyarrow"}, + duckdb_kwargs={"eager_backend": "pyarrow"}, ) ) @@ -32,7 +32,7 @@ def test_collect_kwargs(constructor: Constructor) -> None: @pytest.mark.parametrize( - ("return_type", "expected_cls"), + ("eager_backend", "expected_cls"), [ ("pyarrow", pa.Table), ("polars", pl.DataFrame), @@ -40,7 +40,7 @@ def test_collect_kwargs(constructor: Constructor) -> None: ], ) def test_collect_duckdb( - return_type: Literal["pyarrow", "polars", "pandas"], expected_cls: type + eager_backend: Literal["pyarrow", "polars", "pandas"], expected_cls: type ) -> None: duckdb = pytest.importorskip("duckdb") @@ -48,7 +48,7 @@ def test_collect_duckdb( df_pl = pl.DataFrame(data) # noqa: F841 df = nw.from_native(duckdb.sql("select * from df_pl")) - result = df.lazy().collect(duckdb_kwargs={"return_type": return_type}).to_native() + result = df.lazy().collect(duckdb_kwargs={"eager_backend": eager_backend}).to_native() assert isinstance(result, expected_cls) @@ -62,8 +62,8 @@ def test_collect_duckdb_raise() -> None: with pytest.raises( ValueError, match=( - "Only the following `return_type`'s are supported: pyarrow, pandas and " + "Only the following `eager_backend`'s are supported: pyarrow, pandas and " "polars. Found 'foo'." ), ): - df.lazy().collect(duckdb_kwargs={"return_type": "foo"}) + df.lazy().collect(duckdb_kwargs={"eager_backend": "foo"})