Uso con Spark

Este documento es una introducción rápida al uso de 🤗 Datasets con Spark, con un enfoque particular en cómo cargar un DataFrame de Spark en un objeto Dataset.

A partir de ahí, tienes acceso rápido a cualquier elemento y puedes usarlo como un cargador de datos para entrenar modelos.

Cargar desde Spark

Un objeto Dataset es un envoltorio de una tabla Arrow, que permite lecturas rápidas desde los arrays del dataset hacia tensores de PyTorch, TensorFlow y JAX. La tabla Arrow se mapea en memoria desde el disco, lo que permite cargar datasets más grandes que la RAM disponible.

Puedes obtener un Dataset desde un DataFrame de Spark usando Dataset.from_spark():

>>> from datasets import Dataset
>>> df = spark.createDataFrame(
...     data=[[1, "Elia"], [2, "Teo"], [3, "Fang"]],
...     columns=["id", "name"],
... )
>>> ds = Dataset.from_spark(df)

Los workers de Spark escriben el conjunto de datos en el disco en un directorio de caché como archivos Arrow, y el Dataset se carga desde allí.

Alternativamente, puede omitir la materialización usando IterableDataset.from_spark(), que devuelve un IterableDataset:

>>> from datasets import IterableDataset
>>> df = spark.createDataFrame(
...     data=[[1, "Elia"], [2, "Teo"], [3, "Fang"]],
...     columns=["id", "name"],
... )
>>> ds = IterableDataset.from_spark(df)
>>> print(next(iter(ds)))
{"id": 1, "name": "Elia"}

Caché

Al usar Dataset.from_spark(), el Dataset resultante se almacena en caché; si se llama a Dataset.from_spark() varias veces sobre el mismo DataFrame, no se volverá a ejecutar el trabajo de Spark que escribe el conjunto de datos como archivos Arrow en el disco.

Puede establecer la ubicación de la caché pasando cache_dir= a Dataset.from_spark(). Asegúrese de usar un disco que esté disponible tanto para sus workers como para su máquina actual (el driver).

En una sesión diferente, un DataFrame de Spark no tiene el mismo hash semántico y se volverá a ejecutar un trabajo de Spark y se almacenará en una nueva caché.

Tipos de características

Si su conjunto de datos está compuesto por imágenes, datos de audio o matrices de N dimensiones, puede especificar el argumento features= en Dataset.from_spark() (o IterableDataset.from_spark()):

>>> from datasets import Dataset, Features, Image, Value
>>> data = [(0, open("image.png", "rb").read())]
>>> df = spark.createDataFrame(data, "idx: int, image: binary")
>>> # Also works if you have arrays
>>> # data = [(0, np.zeros(shape=(32, 32, 3), dtype=np.int32).tolist())]
>>> # df = spark.createDataFrame(data, "idx: int, image: array<array<array<int>>>")
>>> features = Features({"idx": Value("int64"), "image": Image()})
>>> dataset = Dataset.from_spark(df, features=features)
>>> dataset[0]
{'idx': 0, 'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=32x32>}

Puede consultar la documentación de Features para conocer todos los tipos de características disponibles.