Uso con PyArrow

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

Esto resulta especialmente útil, ya que permite operaciones rápidas de copia cero, dado que datasets utiliza PyArrow internamente.

Formato del conjunto de datos

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

Para obtener tablas o arreglos de PyArrow en su lugar, puede establecer el formato del conjunto de datos en pyarrow 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("arrow")
>>> ds[0]       # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a"]]
col_1: [[0]]
>>> ds[:2]      # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
>>> ds["data"]  # pa.array
<pyarrow.lib.ChunkedArray object at 0x1394312a0>
[
  [
    "a",
    "b",
    "c",
    "d"
  ]
]

Esto también funciona para objetos IterableDataset obtenidos, por ejemplo, usando load_dataset(..., streaming=True):

>>> ds = ds.with_format("arrow")
>>> for table in ds.iter(batch_size=2):
...     print(table)
...     break
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]

Procesar datos

Las funciones de PyArrow 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. Puede usar funciones de cálculo de Arrow para procesar un conjunto de datos en Dataset.map() o Dataset.filter():

>>> import pyarrow.compute as pc
>>> 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("arrow")
>>> ds = ds.map(lambda t: t.append_column("col_2", pc.add(t["col_1"], 1)), batched=True)
>>> ds[:2]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
col_2: [[1,1]]
>>> ds = ds.filter(lambda t: pc.equal(t["col_0"], "b"), batched=True)
>>> ds[0]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["b"]]
col_1: [[0]]
col_2: [[1]]

We use batched=True because it is faster to process batches of data in PyArrow rather than row by row. It’s also possible to use batch_size= in map() to set the size of each table.

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

Importar o exportar desde PyArrow

Un Dataset es un envoltorio de una tabla de PyArrow; puedes instanciar un Dataset directamente desde la tabla:

ds = Dataset(table)

Puede acceder a la tabla PyArrow de un conjunto de datos mediante Dataset.data, que devuelve una MemoryMappedTable o una InMemoryTable o una ConcatenationTable, dependiendo del origen de los datos de Arrow y de las operaciones que se hayan aplicado.

Esos objetos envuelven la tabla PyArrow subyacente, accesible en Dataset.data.table. Esta tabla contiene todos los datos del conjunto de datos, pero también puede existir un mapeo de índices en Dataset._indices que asocia los índices de las filas del conjunto de datos con los índices de las filas de la tabla PyArrow. Esto puede ocurrir si el conjunto de datos se ha barajado con Dataset.shuffle() o si solo se utiliza un subconjunto de las filas (por ejemplo, después de un Dataset.select()).

En el caso general, puede exportar un conjunto de datos a una tabla de PyArrow usando table = ds.with_format("arrow")[:].