Skip to content

Commit

Permalink
fix(python): Fix regression in scan_parquet (#8071)
Browse files Browse the repository at this point in the history
  • Loading branch information
stinodego committed Apr 9, 2023
1 parent 47dccd9 commit 68c6ea8
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 7 deletions.
10 changes: 7 additions & 3 deletions py-polars/polars/io/ipc/anonymous_scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,19 +28,23 @@ def _scan_ipc_fsspec(


def _scan_ipc_impl( # noqa: D417
source: str, columns: list[str] | None, **kwargs: Any
source: str,
columns: list[str] | None,
predicate: str | None,
n_rows: int | None,
**kwargs: Any,
) -> DataFrame:
"""
Take the projected columns and materialize an arrow table.
Parameters
----------
uri
source
Source URI
columns
Columns that are projected
"""
import polars as pl

return pl.read_ipc(source, columns=columns, **kwargs)
return pl.read_ipc(source, columns=columns, n_rows=n_rows, **kwargs)
10 changes: 7 additions & 3 deletions py-polars/polars/io/parquet/anonymous_scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,19 +28,23 @@ def _scan_parquet_fsspec(


def _scan_parquet_impl( # noqa: D417
uri: str, columns: list[str] | None, **kwargs: Any
source: str,
columns: list[str] | None,
predicate: str | None,
n_rows: int | None,
**kwargs: Any,
) -> DataFrame:
"""
Take the projected columns and materialize an arrow table.
Parameters
----------
uri
source
Source URI
columns
Columns that are projected
"""
import polars as pl

return pl.read_parquet(uri, columns=columns, **kwargs)
return pl.read_parquet(source, columns=columns, n_rows=n_rows, **kwargs)
2 changes: 1 addition & 1 deletion py-polars/polars/lazyframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -502,7 +502,7 @@ def _scan_python_function(
self = cls.__new__(cls)
if isinstance(schema, dict):
self._ldf = PyLazyFrame.scan_from_python_function_pl_schema(
[(name, dt) for name, dt in schema.items()], scan_fn, pyarrow
list(schema.items()), scan_fn, pyarrow
)
else:
self._ldf = PyLazyFrame.scan_from_python_function_arrow_schema(
Expand Down

0 comments on commit 68c6ea8

Please sign in to comment.