Use com Pandas

Este documento é uma introdução rápida ao uso de datasets com Pandas, com foco especial em como processar datasets usando funções do Pandas e como converter um dataset para o Pandas ou a partir do Pandas.

Isso é particularmente útil, pois permite operações rápidas, já que datasets usa PyArrow por baixo dos panos e o PyArrow é bem integrado ao Pandas.

Formato do conjunto de dados

Por padrão, os conjuntos de dados retornam objetos Python comuns: inteiros, números de ponto flutuante, strings, listas, etc.

Para obter DataFrames ou Series do Pandas, você pode definir o formato do dataset como pandas usando 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

Isso também funciona para objetos IterableDataset obtidos, por exemplo, usando 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

Processar dados

As funções do Pandas são geralmente mais rápidas do que funções python escritas manualmente e, portanto, são uma boa opção para otimizar o processamento de dados. Você pode usar funções do Pandas para processar um conjunto de dados em 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

Usamos batched=True porque é mais rápido processar lotes de dados no Pandas do que linha por linha. Também é possível usar batch_size= em map() para definir o tamanho de cada df.

Isso também funciona para IterableDataset.map() e IterableDataset.filter().

Importar ou Exportar do Pandas

Para importar dados do Pandas, você pode usar Dataset.from_pandas():

ds = Dataset.from_pandas(df)

E você pode usar Dataset.to_pandas() para exportar um Dataset para um DataFrame do Pandas:

df = Dataset.to_pandas()