Usar com Spark

Este documento é uma introdução rápida ao uso do 🤗 Datasets com Spark, com foco especial em como carregar um DataFrame do Spark em um objeto Dataset.

A partir daí, você tem acesso rápido a qualquer elemento e pode usá-lo como um carregador de dados para treinar modelos.

Carregar do Spark

Um objeto Dataset é um wrapper de uma tabela Arrow, que permite leituras rápidas de arrays no dataset para tensores do PyTorch, TensorFlow e JAX. A tabela Arrow é mapeada na memória a partir do disco, o que permite carregar datasets maiores que a sua RAM disponível.

Você pode obter um Dataset a partir de um DataFrame do 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)

Os workers do Spark gravam o conjunto de dados em um diretório de cache no disco como arquivos Arrow, e o Dataset é carregado a partir de lá.

Alternativamente, você pode pular a materialização usando IterableDataset.from_spark(), que retorna um 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"}

Cache

Ao usar Dataset.from_spark(), o Dataset resultante é armazenado em cache; se você chamar Dataset.from_spark() várias vezes no mesmo DataFrame, o trabalho do Spark que grava o dataset como arquivos Arrow no disco não será reexecutado.

Você pode definir o local do cache passando cache_dir= para Dataset.from_spark(). Certifique-se de usar um disco que esteja disponível tanto para seus workers quanto para sua máquina atual (o driver).

Em uma sessão diferente, um DataFrame do Spark não possui o mesmo hash semântico e reexecutará um trabalho do Spark, armazenando-o em um novo cache.

Tipos de recursos

Se o seu conjunto de dados é composto por imagens, dados de áudio ou arrays N-dimensionais, você pode especificar o argumento features= em 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>}

Você pode consultar a documentação de Features para conhecer todos os tipos de recursos disponíveis.