Sử dụng với Pandas

Tài liệu này là phần giới thiệu nhanh về việc sử dụng datasets với Pandas, tập trung đặc biệt vào cách xử lý datasets bằng các hàm Pandas, và cách chuyển đổi một dataset sang Pandas hoặc từ Pandas.

Điều này đặc biệt hữu ích vì nó cho phép các thao tác nhanh chóng, do datasets sử dụng PyArrow bên trong và PyArrow được tích hợp tốt với Pandas.

Định dạng tập dữ liệu

Mặc định, các tập dữ liệu trả về các đối tượng Python thông thường: số nguyên, số thực, chuỗi, danh sách, v.v.

Để nhận các DataFrame hoặc Series của Pandas thay thế, bạn có thể đặt định dạng của dataset thành pandas bằng cách sử dụng Dataset.with_format():

>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("pandas")
>>> ds[0]       # pd.DataFrame
  col_0  col_1
0     a    0.0
>>> ds[:2]      # pd.DataFrame
  col_0  col_1
0     a    0.0
1     b    0.0
>>> ds["data"]  # pd.Series
0    a
1    b
2    c
3    d
Name: col_0, dtype: object

Điều này cũng hoạt động với các đối tượng IterableDataset được lấy, ví dụ bằng cách sử dụng load_dataset(..., streaming=True):

>>> ds = ds.with_format("pandas")
>>> for df in ds.iter(batch_size=2):
...     print(df)
...     break
  col_0  col_1
0     a    0.0
1     b    0.0

Xử lý dữ liệu

Các hàm Pandas thường nhanh hơn so với các hàm python viết tay thông thường, do đó chúng là một lựa chọn tốt để tối ưu hóa xử lý dữ liệu. Bạn có thể sử dụng các hàm Pandas để xử lý một tập dữ liệu trong Dataset.map() hoặc Dataset.filter():

>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("pandas")
>>> ds = ds.map(lambda df: df.assign(col_2=df.col_1 + 1), batched=True)
>>> ds[:2]
  col_0  col_1  col_2
0     a    0.0    1.0
1     b    0.0    1.0
>>> ds = ds.filter(lambda df: df.col_0 == "b", batched=True)
>>> ds[0]
  col_0  col_1  col_2
0     b    0.0    1.0

Chúng tôi sử dụng batched=True vì việc xử lý các lô dữ liệu trong Pandas nhanh hơn so với xử lý từng hàng. Bạn cũng có thể sử dụng batch_size= trong map() để đặt kích thước của mỗi df.

Điều này cũng hoạt động với IterableDataset.map() và IterableDataset.filter().

Nhập hoặc Xuất từ Pandas

Để nhập dữ liệu từ Pandas, bạn có thể sử dụng Dataset.from_pandas():

ds = Dataset.from_pandas(df)

Và bạn có thể sử dụng Dataset.to_pandas() để xuất một Dataset sang một Pandas DataFrame:

df = Dataset.to_pandas()