Использование с Pandas

Этот документ представляет собой краткое введение к использованию datasets с Pandas, с особым акцентом на то, как обрабатывать наборы данных с помощью функций Pandas, а также как конвертировать набор данных в 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

Это также работает для объектов IterableDataset, полученных, например, с помощью 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

Обработка данных

Функции 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 быстрее при работе с батчами, чем построчно. Также можно использовать batch_size= в map(), чтобы задать размер каждого 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()