Dask

Dask — это библиотека для параллельных и распределённых вычислений, расширяющая возможности существующей экосистемы Python и PyData.

В частности, мы можем использовать Dask DataFrame для масштабирования рабочих процессов pandas. Dask DataFrame параллелизует pandas для обработки больших табличных данных. Он тесно повторяет API pandas, что упрощает переход от тестирования на одном наборе данных к обработке полного набора. Dask особенно эффективен с Parquet, форматом по умолчанию в Hugging Face Datasets, поскольку он поддерживает богатые типы данных, эффективную фильтрацию по столбцам и сжатие.

Хороший практический пример использования Dask — это распределённая обработка данных или вывод модели на наборе данных. См., например, отличную статью Coiled на тему Scaling AI-Based Data Processing with Hugging Face + Dask.

Чтение и запись

Поскольку Dask использует fsspec для чтения и записи удалённых данных, вы можете использовать пути Hugging Face (hf://) для чтения и записи данных на Hub.

Сначала вам нужно войти в свой аккаунт Hugging Face, например, с помощью:

hf auth login

Затем вы можете создать репозиторий датасета, например, используя:

from huggingface_hub import HfApi

HfApi().create_repo(repo_id="username/my_dataset", repo_type="dataset")

Наконец, в Dask можно использовать пути Hugging Face. Dask DataFrame поддерживает распределённую запись в Parquet на Hugging Face, что использует коммиты для отслеживания изменений в наборе данных:

import dask.dataframe as dd

df.to_parquet("hf://datasets/username/my_dataset")

# or write in separate directories if the dataset has train/validation/test splits
df_train.to_parquet("hf://datasets/username/my_dataset/train")
df_valid.to_parquet("hf://datasets/username/my_dataset/validation")
df_test .to_parquet("hf://datasets/username/my_dataset/test")

Поскольку при этом создаётся отдельный коммит для каждого файла, рекомендуется объединить историю после загрузки:

from huggingface_hub import HfApi

HfApi().super_squash_history(repo_id=repo_id, repo_type="dataset")

Это создаёт репозиторий датасета username/my_dataset, содержащий ваш датасет Dask в формате Parquet. Позже вы можете загрузить его снова:

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/username/my_dataset")

# or read from separate directories if the dataset has train/validation/test splits
df_train = dd.read_parquet("hf://datasets/username/my_dataset/train")
df_valid = dd.read_parquet("hf://datasets/username/my_dataset/validation")
df_test  = dd.read_parquet("hf://datasets/username/my_dataset/test")

Для получения дополнительной информации о путях Hugging Face и их реализации обратитесь к документации клиентской библиотеки по HfFileSystem.

Обработка данных

Чтобы обрабатывать набор данных параллельно с помощью Dask, сначала определите функцию обработки данных для pandas DataFrame или Series, а затем используйте функцию Dask map_partitions, чтобы применить эту функцию ко всем партициям набора данных параллельно:

def dummy_count_words(texts):
    return pd.Series([len(text.split(" ")) for text in texts])

или аналогичная функция с использованием строковых методов pandas (быстрее):

def dummy_count_words(texts):
    return texts.str.count(" ")

В pandas вы можете использовать эту функцию для текстовой колонки:

# pandas API
df["num_words"] = dummy_count_words(df.text)

А в Dask вы можете выполнить эту функцию на каждом разделе:

# Dask API: run the function on every partition
df["num_words"] = df.text.map_partitions(dummy_count_words, meta=int)

Обратите внимание, что вам также необходимо указать meta, который представляет собой тип pandas Series или DataFrame в выводе вашей функции. Это необходимо, так как Dask DataFrame использует ленивый API. Поскольку Dask выполнит обработку данных только после вызова .compute(), ему нужен аргумент meta, чтобы знать тип новой колонки в промежуточное время.

Предикатная и проекционная оптимизация

При чтении данных Parquet из Hugging Face Dask автоматически использует метаданные в файлах Parquet, чтобы пропускать целые файлы или группы строк, если они не нужны. Например, если вы применяете фильтр (предикат) к набору данных Hugging Face в формате Parquet или выбираете подмножество столбцов (проекция), Dask прочитает метаданные файлов Parquet, чтобы отбросить ненужные части, не скачивая их.

Это стало возможным благодаря переосуществлению API Dask DataFrame для поддержки оптимизации запросов, что делает Dask быстрее и надёжнее.

Например, этот поднабор FineWeb-Edu содержит множество файлов Parquet. Если вы можете отфильтровать датасет так, чтобы оставить только текст из недавних дампов CC, Dask пропустит большинство файлов и скачает только те данные, которые соответствуют фильтру:

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/*.parquet")

# Dask will skip the files or row groups that don't
# match the query without downloading them.
df = df[df.dump >= "CC-MAIN-2023"]

Dask также будет читать только те столбцы, которые необходимы для вашего вычисления, и пропускать остальные. Например, если вы удаляете столбец на позднем этапе кода, система не будет загружать его на ранних этапах конвейера, если он не нужен. Это полезно, когда вы хотите работать с подмножеством столбцов или для аналитики:

# Dask will download the 'dump' and 'token_count' needed
# for the filtering and computation and skip the other columns.
df.token_count.mean().compute()

Клиент

Большинство функций в dask оптимизированы для запуска параллельных вычислений на кластере или с помощью локального Client:

import dask.dataframe as dd
from distributed import Client

if __name__ == "__main__":  # needed for creating new processes
    client = Client()
    df = dd.read_parquet(...)
    ...

Для локального использования Client по умолчанию использует Dask LocalCluster с мультипроцессингом. Вы можете вручную настроить мультипроцессинг LocalCluster с

from dask.distributed import Client, LocalCluster
cluster = LocalCluster(n_workers=8, threads_per_worker=8)
client = Client(cluster)

Обратите внимание, что при использовании локально стандартного потокового планировщика без Client DataFrame может работать медленнее после некоторых операций (подробнее здесь).

Подробнее о настройке локального или облачного кластера см. в документации по развертыванию Dask.