搭配 Pandas 使用

本文档是关于如何将 datasets 与 Pandas 结合使用的快速入门指南，重点介绍如何使用 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() 将 Dataset 导出为 Pandas DataFrame

df = Dataset.to_pandas()

在 GitHub 上更新