Hugging Face Datasets 怎样用 flatten_indices 恢复顺序访问性能

如果一个 Hugging Face Datasets 的 Dataset 在 shuffle()、select() 或某些筛选操作后访问变慢，常见原因不是列数据突然变大，而是对象多了一层 indices mapping。读取一行时，程序要先通过映射找到底层 Arrow 表中的真实行，原本可以连续读取的数据块也可能被打散。

解决办法是接收 flatten_indices() 返回的新对象：它会把当前逻辑顺序的数据重写到新的缓存中，移除间接索引，让后续访问回到连续数据路径。官方资料：https://huggingface.co/docs/datasets/process。

为什么 shuffle 之后访问会变慢

Dataset 由 Arrow 表支撑。没有索引映射时，按相邻行读取可以利用连续数据块；而 shuffle() 通常保存一张重新排列后的行号表。调用方看到的顺序是对的，但每次取值都要多一次映射，连续读取也更难保持。

因此，flatten_indices() 不是“把数据重新排序一次的装饰方法”，而是把当前逻辑顺序物化成一份新的连续数据。它会产生写入成本，换来后续训练或批量遍历更简单的读取路径。

最小写法：把返回值赋回 Dataset

所有这类处理方法都返回新的 Dataset，不会把原变量原地改好。因此关键写法是把 flatten_indices() 的返回值重新赋给变量：

from datasets import load_dataset

# 载入一个训练集分片，实际项目中可替换为自己的数据集名称
dataset = load_dataset("nyu-mll/glue", "mrpc", split="train")

# shuffle 产生逻辑上的新顺序，也可能引入 indices mapping
shuffled = dataset.shuffle(seed=42)

# 把当前顺序重写成连续数据；必须接收返回的新 Dataset
flattened = shuffled.flatten_indices()

# 后续训练或批量读取使用重写后的对象
for row in flattened.select(range(3)):
    print(row["label"], row["sentence1"][:30])

这里的 flattened 保留的是 shuffled 对外可见的行顺序，但底层不再依赖那张索引映射。若写成 shuffled.flatten_indices() 却继续使用 shuffled，重写动作本身可能已经完成，后续读取却仍然走旧对象，这是最容易忽略的错误。

flatten_indices 的参数怎么选

默认调用适合先验证流程，但大数据集需要明确写入策略。官方 API 提供的参数可以按下面的职责理解：

例如，希望把重写结果落到指定缓存文件，并使用较大的写入批量，可以这样写：

from pathlib import Path

# 目录应位于有足够空间且生命周期明确的缓存盘
cache_path = Path("./dataset-cache/mrpc-shuffled.arrow")
cache_path.parent.mkdir(parents=True, exist_ok=True)

# 指定缓存位置；num_proc 需要结合 CPU、磁盘吞吐和数据规模调整
flattened = shuffled.flatten_indices(
    cache_file_name=str(cache_path),
    writer_batch_size=1000,
    num_proc=2,
)

# 后续把 flattened 传给训练器，而不是继续使用 shuffled
print(flattened.column_names)

writer_batch_size 不是“越大越快”的开关。它影响写入器的批量和临时内存，应该在真实机器的内存、磁盘和并发条件下调整；num_proc 也不是免费加速，多个进程可能让共享存储成为瓶颈。

什么时候不应该马上 flatten

flatten_indices() 会重写整个逻辑数据集，因此它适合“后面要反复读取”的场景，不适合每个小实验都调用一次。下面三种情况可以先不做：

• 只抽取很小的样本做一次性检查，直接使用当前对象更省事。
• 数据集即将再次筛选、打乱或拆分，过早重写会产生一次额外的磁盘写入。
• 训练流程接受流式遍历且不要求精确的全量随机访问，可以考虑 to_iterable_dataset()。

官方文档还给出另一条思路：先把数据转换为 IterableDataset，再使用带缓冲区的近似 shuffle。它避免了把所有随机顺序一次性物化，代价是访问模型从可按索引取行变成惰性迭代，不能把两种对象当成完全等价的替换。

常见误区与排查顺序

1. 没有接收返回值：确认训练器、DataLoader 或后续 map() 使用的是重写后的变量。
2. 把重写当成原地操作：先把新对象命名为 flattened，等流程稳定后再覆盖旧变量，便于排查顺序问题。
3. 忽略磁盘空间：重写期间要同时考虑原始数据、缓存文件和临时分片，不要只看数据集的逻辑大小。
4. 只盯着 CPU：重写与读取速度还受磁盘连续写入、缓存目录和进程数影响，增加 num_proc 前先看存储吞吐。
5. 把近似 shuffle 当精确 shuffle：IterableDataset.shuffle() 使用缓冲区，是另一种随机化模型；如果任务依赖可复现的精确行顺序，要保留明确的 seed 与数据处理链。

结论与相关问题

一句话概括：如果 shuffle 或筛选后的 Dataset 要被反复顺序读取，就把 flatten_indices() 当成一次“用写入换读取”的整理步骤，并接收它返回的新对象；如果只是短暂实验或适合流式遍历，则先评估是否需要这次全量重写。

flatten_indices 会改变 shuffle 后的顺序吗？

它的目标是移除索引映射并按当前逻辑顺序重写数据，不是重新随机一次。文章示例中的顺序仍由前面的 shuffle(seed=42) 决定。

为什么指定了 cache_file_name 还要保留原始数据？

缓存文件是重写结果，原始数据仍可能被其他流程使用。是否清理原始数据要由你的缓存生命周期和可恢复需求决定，不能把缓存路径当作自动迁移或自动删除开关。

num_proc 应该设置成 CPU 核数吗？

不一定。它还受存储吞吐、数据格式和机器内存约束，建议从较小值开始观察整体重写耗时与资源占用。