Использование со Spark

Этот документ представляет собой краткое введение к использованию 🤗 Datasets с Spark, с особым акцентом на то, как загрузить Spark DataFrame в объект Dataset.

Оттуда у вас есть быстрый доступ к любому элементу, и вы можете использовать его как загрузчик данных для обучения моделей.

Загрузка из Spark

Объект Dataset представляет собой обёртку над таблицей Arrow, которая позволяет быстро считывать данные из массивов набора в тензоры PyTorch, TensorFlow и JAX. Таблица Arrow отображается в память с диска (memory-mapped), что позволяет загружать наборы данных, превышающие объём доступной оперативной памяти.

Вы можете получить Dataset из Spark DataFrame с помощью 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)

Рабочие процессы Spark записывают набор данных на диск в каталог кэша в формате файлов Arrow, а Dataset загружается оттуда.

В качестве альтернативы можно пропустить материализацию, используя IterableDataset.from_spark(), который возвращает 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"}

Кэширование

При использовании Dataset.from_spark() результирующий Dataset кэшируется; если вы несколько раз вызываете Dataset.from_spark() для одного и того же DataFrame, Spark-задача, записывающая датасет в виде файлов Arrow на диск, не будет выполняться повторно.

Вы можете указать расположение кэша, передав cache_dir= в Dataset.from_spark(). Убедитесь, что используете диск, доступный как вашим воркерам, так и текущей машине (драйверу).

В другой сессии у Spark DataFrame будет другой семантический хеш, поэтому Spark-задача будет выполнена заново, а её результат сохранён в новый кэш.

Типы признаков

Если ваш набор данных состоит из изображений, аудиоданных или N-мерных массивов, вы можете указать аргумент features= в Dataset.from_spark() (или 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>}

Вы можете ознакомиться с документацией Features, чтобы узнать обо всех доступных типах признаков.