与 Pandas 配合使用

本文档简要介绍了如何使用 Pandas 处理 datasets，重点说明如何使用 Pandas 函数处理数据集，以及如何在数据集与 Pandas 之间进行转换。

这一点尤其有用，因为它可以实现快速操作，因为 datasets 底层使用 PyArrow，而 PyArrow 与 Pandas 集成良好。

数据集格式

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

To get Pandas DataFrames or Series instead, you can set the format of the dataset to pandas using 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 函数更快，因此是优化数据处理的理想选择。您可以在 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)

And you can use Dataset.to_pandas() to export a Dataset to a Pandas DataFrame:

df = Dataset.to_pandas()

Update on GitHub

Datasets

与 Pandas 配合使用

数据集格式

处理数据

从 Pandas 导入或导出