Utiliser avec Spark

Ce document est une introduction rapide à l'utilisation de 🤗 Datasets avec Spark, avec un accent particulier sur la manière de charger un DataFrame Spark dans un objet Dataset.

À partir de là, vous avez un accès rapide à n'importe quel élément et vous pouvez l'utiliser comme chargeur de données pour entraîner des modèles.

Charger depuis Spark

Un objet Dataset est un conteneur pour un tableau Arrow, qui permet des lectures rapides depuis les tableaux du jeu de données vers des tenseurs PyTorch, TensorFlow et JAX. Le tableau Arrow est mappé en mémoire depuis le disque, ce qui permet de charger des jeux de données plus volumineux que la RAM disponible.

Vous pouvez obtenir un Dataset à partir d'un DataFrame Spark en utilisant 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)

Les workers Spark écrivent le jeu de données sur le disque dans un répertoire de cache sous forme de fichiers Arrow, et le Dataset est chargé depuis ce répertoire.

Sinon, vous pouvez éviter la matérialisation en utilisant IterableDataset.from_spark(), qui renvoie 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"}

Mise en cache

Lors de l’utilisation de Dataset.from_spark(), le Dataset résultant est mis en cache ; si vous appelez Dataset.from_spark() plusieurs fois sur le même DataFrame, le travail Spark qui écrit le jeu de données sous forme de fichiers Arrow sur le disque ne sera pas relancé.

Vous pouvez définir l'emplacement du cache en passant cache_dir= à Dataset.from_spark(). Assurez-vous d'utiliser un disque accessible à la fois par vos workers et par votre machine actuelle (le driver).

Dans une session différente, un DataFrame Spark ne possède pas le même hash sémantique et relancera un travail Spark pour le stocker dans un nouveau cache.

Types de fonctionnalités

Si votre jeu de données est composé d’images, de données audio ou de tableaux N-dimensionnels, vous pouvez spécifier l’argument features= dans Dataset.from_spark() (ou 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>}

Vous pouvez consulter la documentation Features pour connaître tous les types de fonctionnalités disponibles.