Skip to content

Commit

Permalink
remove kwargs in collect method
Browse files Browse the repository at this point in the history
  • Loading branch information
jackxxu committed Nov 3, 2024
1 parent ff10b38 commit 4679380
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 3 deletions.
6 changes: 3 additions & 3 deletions py-polars/polars/lazyframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -1830,10 +1830,11 @@ def collect(
collapse_joins: bool = True,
no_optimization: bool = False,
streaming: bool = False,
new_streaming: bool = False,
engine: EngineType = "cpu",
background: bool = False,
_eager: bool = False,
**_kwargs: Any,
post_opt_callback: Callable[..., Any] | None = None,
) -> DataFrame | InProcessQuery:
"""
Materialize this LazyFrame into a DataFrame.
Expand Down Expand Up @@ -1982,7 +1983,6 @@ def collect(
│ c ┆ 6 ┆ 1 │
└─────┴─────┴─────┘
"""
new_streaming = _kwargs.get("new_streaming", False)

if no_optimization or _eager:
predicate_pushdown = False
Expand Down Expand Up @@ -2046,7 +2046,7 @@ def collect(
engine = GPUEngine()
callback = partial(cudf_polars.execute_with_cudf, config=engine)
# Only for testing purposes
callback = _kwargs.get("post_opt_callback", callback)
callback = post_opt_callback or callback
return wrap_df(ldf.collect(callback))

@overload
Expand Down
5 changes: 5 additions & 0 deletions py-polars/tests/unit/lazyframe/test_lazyframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -907,6 +907,11 @@ def test_collect_all(df: pl.DataFrame, no_optimization: bool) -> None:
assert cast(float, out[1].item()) == 12.0


def test_collect_unexpected_kwargs(df: pl.DataFrame) -> None:
with pytest.raises(TypeError, match="unexpected keyword argument"):
df.lazy().collect(common_subexpr_elim=False)


def test_spearman_corr() -> None:
ldf = pl.LazyFrame(
{
Expand Down

0 comments on commit 4679380

Please sign in to comment.