Uso con Pandas

Este documento es una introducción rápida al uso de datasets con Pandas, con un enfoque particular en cómo procesar conjuntos de datos mediante funciones de Pandas y cómo convertir un conjunto de datos a Pandas o desde Pandas.

Esto resulta especialmente útil, ya que permite operaciones rápidas, dado que datasets utiliza PyArrow internamente y PyArrow está bien integrado con Pandas.

Formato del conjunto de datos

Por defecto, los conjuntos de datos devuelven objetos de Python estándar: enteros, flotantes, cadenas, listas, etc.

Para obtener DataFrames o Series de Pandas en su lugar, puede establecer el formato del conjunto de datos en 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

Esto también funciona para objetos IterableDataset obtenidos, por ejemplo, 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

Procesar datos

Las funciones de Pandas son generalmente más rápidas que las funciones de python escritas a mano, por lo que son una buena opción para optimizar el procesamiento de datos. Puedes usar funciones de Pandas para procesar un conjunto de datos en Dataset.map() o 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 es más rápido procesar lotes de datos en Pandas que fila por fila. También es posible usar batch_size= en map() para establecer el tamaño de cada df.

Esto también funciona para IterableDataset.map() y IterableDataset.filter().

Importar o exportar desde Pandas

Para importar datos desde Pandas, puedes usar Dataset.from_pandas():

ds = Dataset.from_pandas(df)

Y puedes usar Dataset.to_pandas() para exportar un Dataset a un DataFrame de Pandas:

df = Dataset.to_pandas()