Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Capture compression and open binary if present. #419

Merged
merged 2 commits into from
Nov 14, 2024
Merged

Capture compression and open binary if present. #419

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e069f4b
Capture compression and open binary if present.
delucchi-cmu Nov 14, 2024
b3c1d38
Add non-utf8 encoded test
delucchi-cmu Nov 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/hats/io/file_io/file_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,18 +107,21 @@ def load_csv_to_pandas(file_pointer: str | Path | UPath, **kwargs) -> pd.DataFra


def load_csv_to_pandas_generator(
file_pointer: str | Path | UPath, chunksize=10_000, **kwargs
file_pointer: str | Path | UPath, *, chunksize=10_000, open_mode=None, compression=None, **kwargs
) -> Generator[pd.DataFrame]:
"""Load a csv file to a pandas dataframe
Args:
file_pointer: location of csv file to load
file_system: fsspec or pyarrow filesystem, default None
chunksize (int): number of rows to load per chunk
compression (str): for compressed CSVs, the manner of compression. e.g. 'gz', 'bzip'.
**kwargs: arguments to pass to pandas `read_csv` loading method
Returns:
pandas dataframe loaded from CSV
"""
file_pointer = get_upath(file_pointer)
with file_pointer.open("r", **kwargs) as csv_file:
if open_mode is None:
open_mode = "r" if compression is None else "rb"
delucchi-cmu marked this conversation as resolved.
Show resolved Hide resolved
with file_pointer.open(mode=open_mode, compression=compression, **kwargs) as csv_file:
with pd.read_csv(csv_file, chunksize=chunksize, **kwargs) as reader:
yield from reader

Expand Down