与 Pandas 配合使用

本文档简要介绍了如何将 `datasets` 与 Pandas 结合使用，特别侧重于如何使用 Pandas 函数处理数据集，以及如何将数据集与 Pandas 互相转换。

这特别有用，因为它允许快速操作，因为 `datasets` 底层使用 PyArrow，而 PyArrow 与 Pandas 结合得很好。

数据集格式

默认情况下，数据集返回常规的Python对象：整数、浮点数、字符串、列表等。

要获取 Pandas DataFrame 或 Series，您可以使用 Dataset.with_format() 将数据集的格式设置为 `pandas`。

>>> 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

这也适用于例如使用 `load_dataset(..., streaming=True)` 获取的 `IterableDataset` 对象。

>>> 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

处理数据

Pandas 函数通常比普通的手写 Python 函数更快，因此它们是优化数据处理的好选择。您可以在 Dataset.map() 或 Dataset.filter() 中使用 Pandas 函数来处理数据集。

>>> 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

我们使用 `batched=True`，因为在 Pandas 中批量处理数据比逐行处理更快。也可以在 `map()` 中使用 `batch_size=` 来设置每个 `df` 的大小。

这也适用于 IterableDataset.map() 和 IterableDataset.filter()。

从 Pandas 导入或导出

要从 Pandas 导入数据，您可以使用 Dataset.from_pandas()

ds = Dataset.from_pandas(df)

您可以使用 Dataset.to_pandas() 将数据集导出到 Pandas DataFrame

df = Dataset.to_pandas()

< > 在 GitHub 上更新