Utiliser avec Pandas

Ce document est une introduction rapide à l’utilisation de datasets avec Pandas, en mettant particulièrement l’accent sur la manière de traiter des datasets à l’aide de fonctions Pandas, ainsi que sur la conversion d’un dataset vers Pandas ou depuis Pandas.

C'est particulièrement utile car cela permet des opérations rapides, datasets utilisant PyArrow en arrière-plan et PyArrow étant bien intégré avec Pandas.

Format du jeu de données

Par défaut, les jeux de données renvoient des objets Python classiques : des entiers, des flottants, des chaînes de caractères, des listes, etc.

Pour obtenir des DataFrames ou des Series Pandas à la place, vous pouvez définir le format du jeu de données sur pandas à l’aide de 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

Cela fonctionne également pour les objets IterableDataset obtenus, par exemple, à l'aide de 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

Traiter les données

Les fonctions Pandas sont généralement plus rapides que les fonctions Python écrites à la main et constituent donc une bonne option pour optimiser le traitement des données. Vous pouvez utiliser les fonctions Pandas pour traiter un jeu de données dans Dataset.map() ou 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

Nous utilisons batched=True car il est plus rapide de traiter des lots de données dans Pandas que ligne par ligne. Il est également possible d’utiliser batch_size= dans map() pour définir la taille de chaque df.

Cela fonctionne également pour IterableDataset.map() et IterableDataset.filter().

Importer ou exporter depuis Pandas

Pour importer des données depuis Pandas, vous pouvez utiliser Dataset.from_pandas() :

ds = Dataset.from_pandas(df)

Vous pouvez utiliser Dataset.to_pandas() pour exporter un Dataset vers un DataFrame Pandas :

df = Dataset.to_pandas()