与 Pandas 一起使用

本文档快速介绍如何将 datasets 与 Pandas 一起使用，特别关注如何使用 Pandas 函数处理数据集，以及如何将数据集转换为 Pandas 或从 Pandas 转换。

这非常有用，因为它允许快速操作，因为 datasets 在底层使用 PyArrow，而 PyArrow 与 Pandas 很好地集成。

数据集格式

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

要获取 Pandas DataFrame 或 Series，您可以将数据集的格式设置为 pandas，使用 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

这也适用于通过 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 函数更快，因此它们是优化数据处理的良好选择。您可以使用 Pandas 函数在 Dataset.map() 或 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

我们使用 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 上更新