# Dask

[Dask](https://www.dask.org/?utm_source=hf-docs) 是一个并行和分布式计算库，可扩展现有的 Python 和 PyData 生态系统。

具体而言，我们可以使用 [Dask DataFrame](https://docs.dask.org/en/stable/dataframe.html?utm_source=hf-docs) 来扩展 pandas 工作流。Dask DataFrame 通过并行化 pandas 来处理大型表格数据。它紧密贴合 pandas API，使得从在单个数据集上进行测试过渡到处理完整数据集变得十分简单。Dask 在处理 Parquet 时尤为高效，Parquet 是 Hugging Face Datasets 的默认格式，它支持丰富的数据类型、高效的列式过滤以及压缩。

Dask 的一个良好实用场景是以分布式方式在数据集上运行数据处理或模型推理。例如，请参阅 [Coiled](https://www.coiled.io/?utm_source=hf-docs) 关于 [使用 Hugging Face + Dask 扩展基于 AI 的数据处理](https://huggingface.co/blog/dask-scaling) 的优秀博客文章。

## 读取和写入

由于 Dask 使用 [fsspec](https://filesystem-spec.readthedocs.io) 来读写远程数据，您可以使用 Hugging Face 路径（[`hf://`](/docs/huggingface_hub/guides/hf_file_system#integrations)）来读写 Hub 上的数据。

首先，您需要使用 Hugging Face 账户[登录](/docs/huggingface_hub/quick-start#login)，例如使用以下命令：

```
hf auth login
```

然后，您可以[创建数据集仓库](/docs/huggingface_hub/quick-start#create-a-repository)，例如使用：

```python
from huggingface_hub import HfApi

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

最后，你可以在 Dask 中使用 [Hugging Face 路径](/docs/huggingface_hub/guides/hf_file_system#integrations)。Dask DataFrame 支持向 Hugging Face 上的 Parquet 进行分布式写入，该过程使用提交（commits）来跟踪数据集的变更：

```python
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")
```

由于这会为每个文件创建一个提交，建议在上传后压缩历史记录：

```python
from huggingface_hub import HfApi

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

这将创建一个数据集仓库 `username/my_dataset`，其中包含以 Parquet 格式存储的 Dask 数据集。您可以稍后重新加载它：

```python
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 的客户端库文档](/docs/huggingface_hub/guides/hf_file_system)。

## 处理数据

要使用 Dask 并行处理数据集，您可以先为 pandas DataFrame 或 Series 定义数据处理函数，然后使用 Dask 的 `map_partitions` 函数将该函数并行地应用于数据集的所有分区：

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

或一个使用 pandas 字符串方法的类似函数（速度更快）：

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

在 pandas 中，你可以将此函数应用于文本列：

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

在 Dask 中，你可以在每个分区上运行此函数：

```python
# 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` 参数来在此期间了解新列的类型。

## 谓词和投影下推

从 Hugging Face 读取 Parquet 数据时，Dask 会自动利用 Parquet 文件中的元数据，在不需要时跳过整个文件或行组。例如，当您对 Parquet 格式的 Hugging Face Dataset 应用过滤器（谓词），或选择列的子集（投影）时，Dask 会读取 Parquet 文件的元数据，以丢弃不需要的部分，而无需下载它们。

这得益于对 [Dask DataFrame API 的重新实现](https://docs.coiled.io/blog/dask-dataframe-is-fast.html?utm_source=hf-docs)，以支持查询优化，从而使 Dask 更快、更稳健。

例如，FineWeb-Edu 的这个子集包含许多 Parquet 文件。如果你能过滤数据集以保留来自近期 CC 转储的文本，Dask 将跳过大部分文件，仅下载符合过滤条件的数据：

```python
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 还会仅读取计算所需的列，并跳过其余列。例如，如果你在代码的后期删除了某一列，而该列在早期并不需要，那么它就不会在管道的前端加载该列。当你想要操作列的子集或进行数据分析时，这一点非常有用：

```python
# 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` 进行了优化，以启动并行计算：

```python
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` 手动配置多进程：

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

请注意，如果您在本地使用默认的线程调度器且未使用 `Client`，某些操作后 DataFrame 的性能可能会变慢（更多详细信息请参见 [此处](https://github.com/dask/dask-expr/issues/1181)）。

有关设置本地或云集群的更多信息，请参阅 [Dask 部署文档](https://docs.dask.org/en/latest/deploying.html)。