Hugging Face Datasets 导出 Parquet 时怎么控制分片大小

我在把数据集交给下游训练或上传 Hub 时，最先遇到的不是 Parquet 格式本身，而是“分片大小”这个词同时指向了三件事：一次写入多少行、最终生成多少个本地文件，以及远程仓库允许单个文件多大。Hugging Face Datasets 官方文档把这些能力分在不同接口里。

官方文档：https://huggingface.co/docs/datasets/process

如果目标是降低单次写入的内存压力，用 to_parquet(batch_size=...)；如果目标是得到固定数量的本地 Parquet 文件，用 Dataset.shard(num_shards, index) 后逐个导出；如果目标是上传 Hub 时限制单文件大小，用 push_to_hub(max_shard_size=...)，需要固定文件数时再用 num_shards。

• batch_size 是写入批次，不直接决定输出文件个数。
• num_shards 是本地或 Hub 的分片数量，必须结合输出路径或上传接口理解。
• max_shard_size 是 Hub 上传分片的大小上限，适合按容量约束交付文件。

先把“分片大小”拆成三个问题

Dataset.to_parquet() 接收 path_or_buf 和可选的 batch_size。官方 API 对这个参数的解释是：一次加载到内存并写出的批次大小；默认行为会以行组的未压缩字节上限为目标。它解决的是写入过程的内存和行组粒度，并不会因为把 batch_size 改成 5000 就自动创建五个文件。

本地要得到多个文件，需要显式把原数据集切成多个 Dataset。shard(num_shards=n, index=i) 返回第 i 片；默认是连续区间，所有片按索引顺序拼接后仍能还原原顺序。这样“文件数量”由 n 决定，“每次写入多少行”仍由每个分片导出时的 batch_size 决定。

本地固定文件数量：shard 后逐个导出

我需要给离线任务准备四个本地文件时，会把“分成几片”写在外层循环里，再把写入批次作为另一个参数传入。下面的示例用内存数据集演示完整关系，换成 load_dataset() 返回的 Dataset 也一样。

from pathlib import Path

from datasets import Dataset


dataset = Dataset.from_dict({
    "text": ["样本-0", "样本-1", "样本-2", "样本-3", "样本-4", "样本-5"],
    "label": [0, 1, 0, 1, 0, 1],
})

num_shards = 3
output_dir = Path("parquet_parts")
output_dir.mkdir(parents=True, exist_ok=True)

for index in range(num_shards):
    # contiguous=True 让每个分片拿到连续区间，便于按文件顺序合并。
    part = dataset.shard(
        num_shards=num_shards,
        index=index,
        contiguous=True,
    )
    # batch_size 只控制单次加载和写入批次，不改变分片文件数量。
    part.to_parquet(
        output_dir / f"data-{index:05d}.parquet",
        batch_size=1024,
    )

这里会得到三个输出路径。数据行不能被平均整除时，前面的连续分片会多一行，最后的文件可能更短；这不是丢数据，而是 len(dataset) 除以 num_shards 后的余数分配。不要把“每个文件字节数完全相同”当成正确性的判断条件，因为文本长度和嵌套字段会让不同文件的压缩后大小自然不同。

上传 Hub：用容量上限或数量约束

如果数据最终要进入 Hugging Face Hub，直接调用 push_to_hub() 更省事。官方接口提供 max_shard_size 来设置上传分片的最大大小，也提供 num_shards 明确指定要写出的分片数量。当只给出容量上限时，库会根据它估算分片数量；当业务要求固定文件数时，直接给出 num_shards 更容易在目录规划和下游脚本中保持稳定。

from datasets import load_dataset


dataset = load_dataset("your-org/your-source-dataset", split="train")

# 让 Hub 按单文件容量上限估算分片数量；实际上传前需要完成登录授权。
dataset.push_to_hub(
    "your-org/your-output-dataset",
    split="train",
    max_shard_size="512MB",
)

# 如果下游脚本依赖固定数量，可以改为显式指定分片数。
# dataset.push_to_hub("your-org/your-output-dataset", split="train", num_shards=8)

这两个参数回答的是不同的交付约束：容量上限更适合仓库、网络和单文件处理能力有限的场景；固定数量更适合已经约定好并发任务数或目录索引的场景。不要把本地 to_parquet(batch_size=1024) 的数值直接当成 Hub 的 1024 个文件。

我会这样选择参数

常见问题

把 batch_size 调大就会生成更多 Parquet 文件吗？

不会。它控制的是导出时一次加载和写入的批次；要生成多个本地文件，需要显式调用 shard() 并为每片提供不同路径。

数据不能被 num_shards 整除会漏掉最后几行吗？

连续分片会把余数分配给前面的分片，所有分片按原顺序拼接仍覆盖完整数据。实际使用时要保持每个 index 都被导出，不能只处理前几个文件。

本地导出和 push_to_hub 应该选哪个？

需要自定义目录、离线交付或本地检查时先用 shard()+to_parquet()；要直接发布到 Hub 并按容量或数量管理远程文件时用 push_to_hub()。

我最后只记住一个边界：batch_size 管写入节奏，shard 管本地文件拆分，max_shard_size 和 num_shards 管 Hub 交付。先确定约束，再选参数，后面的 Parquet 目录就不会因为“分片大小”这个模糊词而失控。