Hugging Face Datasets 怎样分批导出字典而不一次读取全量

如果直接调用 ds.to_dict()，Hugging Face Datasets 会把整个数据集组织成一个 Python 字典；当数据量较大时，真正需要控制的是“每次转换多少行”。更合适的写法是 ds.to_dict(batch_size=256, batched=True)：它返回一个生成器，每次产出一个批次的列式字典，处理完当前批次后再继续取下一批。

官方地址：https://huggingface.co/docs/datasets/en/package_reference/main_classes

下面用一个小型数据集说明接口语义。重点不是把数据集变成“零内存”对象，而是避免先聚合出一个包含全量列值的 Python 字典。

to_dict 的 batched=True 到底返回什么

Dataset.to_dict 有两个关键参数：batched 为 True 时返回批次生成器，batch_size 控制每个批次的行数；不传 batched=True 时，调用会一次返回完整字典。每个批次仍然是“列名到列表”的列式结构，而不是“每行一个字典”的列表。

from datasets import Dataset

# 构造一个带两列的小数据集，实际项目中也可以替换成 load_dataset 的结果。
ds = Dataset.from_dict({
    "id": [101, 102, 103, 104, 105],
    "text": ["甲", "乙", "丙", "丁", "戊"],
})

# batched=True 返回生成器；batch_size 表示每个批次包含的行数。
batch_stream = ds.to_dict(batch_size=2, batched=True)

for batch_no, batch in enumerate(batch_stream, start=1):
    # batch 是列式字典，例如 {"id": [101, 102], "text": ["甲", "乙"]}。
    first_column = next(iter(batch.values()), [])
    row_count = len(first_column)
    print(f"批次 {batch_no}: {row_count} 行，列={list(batch)}")

上例的最后一批可能少于 batch_size，这是正常结果。不要用固定的 batch_size 推断最后一批行数，应当从实际批次的列列表长度计算。

按批次写出字典文件

如果下游程序需要保留“列名对应一组值”的字典形状，可以每个批次写一个 JSON 文件。文件名带上批次编号，失败时能直接定位到哪一批；写完循环体后没有额外的全量列表累积。

import json
from pathlib import Path
from datasets import load_dataset

# split 只示范接口用法；实际项目按自己的数据集和 split 替换。
ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
out_dir = Path("dataset-batches")
out_dir.mkdir(parents=True, exist_ok=True)

# 每次只从生成器取一个列式字典；不要在外层再用 list(...) 收集所有批次。
for batch_no, batch in enumerate(ds.to_dict(batch_size=256, batched=True), start=1):
    # 用第一列的长度记录本批行数，空数据集也能安全处理。
    first_column = next(iter(batch.values()), [])
    row_count = len(first_column)
    target = out_dir / f"part-{batch_no:05d}.json"

    # ensure_ascii=False 保留中文；indent 只影响文件可读性，不改变批次边界。
    with target.open("w", encoding="utf-8") as fp:
        json.dump(batch, fp, ensure_ascii=False)

    print(f"已写出 {target}，行数={row_count}")

这里的 JSON 是列式批次，例如 {"id": [101, 102], "text": ["甲", "乙"]}。如果接收方要求行式记录，先在当前批次内转换，不要把所有批次再合并成一个大列表。

import json
from pathlib import Path

def write_rows(batch, target: Path):
    # 所有列的列表长度应对应同一批行；按列位置重新组合为行对象。
    keys = list(batch)
    columns = [batch[key] for key in keys]
    with target.open("w", encoding="utf-8") as fp:
        for values in zip(*columns):
            # 每次只构造一行，适合 JSONL 等逐行消费场景。
            fp.write(json.dumps(dict(zip(keys, values)), ensure_ascii=False))
            fp.write("\n")

需要 JSON 文件时直接调用 to_json

如果目标就是 JSON 文件，而不是要在 Python 中检查或改造每个批次，优先考虑 Dataset.to_json(path, batch_size=...)。官方接口会按指定批量加载并写出，调用方不必先调用 to_dict 再手动序列化。

from datasets import load_dataset

# 目标是文件时直接让 Datasets 按批量写出，避免先创建完整 Python 字典。
ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
written = ds.to_json("reviews.jsonl", batch_size=256)

# written 是写入结果的返回值；真正的文件内容仍由目标格式决定。
print(f"写出字节数：{written}")

to_dict(batched=True) 更适合“每批要参与业务处理”的场景，例如清洗、统计、写入多个分区；to_json(batch_size=...) 更适合“只想把数据导成 JSON 文件”的场景。两者都要根据行大小和字段复杂度调整批量，而不是盲目使用很大的数字。

从列式批次恢复成行式记录

很多人看到 batch 后会写成 for row in batch，但字典迭代得到的是列名。正确做法是在当前批次内用列列表的同一位置配成行；前面的 write_rows 已经给出了 JSONL 写法。

如果只想在内存中拿到批量数据，也可以使用 Dataset.iter(batch_size=...)。这个接口的目标是迭代批次，返回值会受到当前格式设置影响；它和 to_dict 的选择，应以“是否要明确得到 Python 字典并自行序列化”为判断标准。

大数据集还要注意流式读取边界

分批导出降低的是“Python 端一次聚合全量结果”的压力，并不意味着磁盘、Arrow 缓存、解码对象和下游写入都不占内存。单行很大、图片字段较多或嵌套结构复杂时，batch_size=256 也可能过大；应先用较小批量观察单批耗时和输出体积，再逐步调大。

如果数据源本身就不适合完整加载，可以研究 load_dataset(..., streaming=True) 或把已加载的 Dataset 转成 to_iterable_dataset()。后者返回可迭代数据集，并支持在迭代时进行过滤和映射；它解决的是迭代方式，不会自动替你设计文件分片、断点和失败重试。

常见问题与排查建议

为什么 to_dict() 还是占用很多内存？

因为默认调用会返回完整数据集的 Python 字典。改为 to_dict(batch_size=256, batched=True)，并确认没有在外层使用 list(generator) 或把各批次追加到全量列表。

最后一批少于 batch_size 是错误吗？

不是。只要没有设置丢弃尾批的逻辑，最后一批自然可以更短；写文件时从实际列长度记录行数即可。

字典批次可以直接当成行列表吗？

不可以。批次是列式字典，遍历字典得到列名；需要行式数据时，在当前批次内用 zip 对齐各列。

什么时候直接使用 to_json？

当目标只是 JSON 文件，且不需要在 Python 中逐批修改字段时，直接调用 to_json(batch_size=...) 更简洁；需要自定义清洗或分区命名时，再选 to_dict(batched=True)。

总结

解决“分批导出字典而不一次读取全量”的关键只有两点：第一，明确使用 to_dict(batched=True) 让结果成为批次生成器；第二，批次用完就写出或交给下游，不要再聚合回全量对象。若目标是纯 JSON 文件，直接使用带 batch_size 的 to_json；若目标是流式处理链，再结合 IterableDataset 评估数据源和下游缓冲。