Datasets map 的 batched 参数怎样改变函数输入

Hugging Face Datasets 的 Dataset.map 默认把一条记录传给回调；设置 batched=True 后，默认 Python 格式下的输入变成“列名对应列表”的批字典。例如单条模式里的 text 是字符串，批量模式里的 batch["text"] 是一组字符串。回调应按批生成列列表，不能继续把它当成单个文本调用 strip。

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

先区分一条记录与一批列列表

本文讨论普通 Dataset、默认格式、未设置 input_columns 的常见用法。假设数据有 text 与 item_id 两列：单条回调读取的是某一行的两个值；批量回调读取的是两列各自的一组值。同一索引位置仍然属于同一条输入记录，因此遍历两列时要保持对应关系。

batched 只改变映射函数的调用粒度，不自动增加多进程，也不等于训练阶段 DataLoader 的 batch。对同一个数据集，逐条和批量写法都可以生成相同的新列，只是函数接收的对象不同。

写一个保持行数不变的批量映射

下面示例只构造三条本地记录，分别演示两种回调。需要环境中已安装 datasets；示例不加载分词器，也不请求远程模型。

from datasets import Dataset

# 只在内存中构造小数据，不下载数据集或模型。
source = Dataset.from_dict({
    "text": ["  red leaf ", "blue sky", "green field"],
    "item_id": [31, 32, 33],
})

def clean_one(example):
    # 单条模式下，text 是一个字符串。
    return {"clean_text": example["text"].strip()}

def clean_batch(batch):
    # 批量模式下，text 是字符串列表；每条记录产生一个结果。
    cleaned = [text.strip() for text in batch["text"]]
    return {"clean_text": cleaned}

single = source.map(clean_one)
batched = source.map(clean_batch, batched=True, batch_size=2)

# 预期两种写法得到相同的新列；source 仍只有原来的两列。
print(single["clean_text"])
print(batched["clean_text"])
print(source.column_names)

按这段代码，两份结果的新列都应为去除两端空白的三条字符串，source.column_names 仍对应 text 和 item_id。这些是根据示例逻辑给出的预期，不是本机运行记录。map 返回新 Dataset，使用时要接住返回对象。

本例没有删除原列，所以 text 和 item_id 继续保留，clean_text 作为新列加入。若返回的键与原列同名，则用于更新该列；批量回调返回的是列内容，不能简单返回“若干字典组成的行列表”并期待它自动当成标准输出批。

扩展行数时处理原有列

batched=True 允许一批 N 条输入生成 M 条输出，例如拆词或切片。但同一输出批内的所有保留列必须有相同的行数。如果新列有更多记录、原列还保持输入数量，Arrow 表就无法成立。处理方式是同步重建需要保留的列，或者在 map 中通过 remove_columns 去掉不再匹配的原列。

# 沿用前面的 source，把每条文本按空白拆成多个词条。
def split_words(batch):
    words = []
    parent_ids = []
    for text, item_id in zip(batch["text"], batch["item_id"]):
        for word in text.split():
            words.append(word)
            # 一条来源拆成多行时，同步复制其来源编号。
            parent_ids.append(item_id)
    return {"word": words, "parent_id": parent_ids}

expanded = source.map(
    split_words,
    batched=True,
    batch_size=2,
    # 删除旧列，避免旧批次的行数与拆分后的新行数冲突。
    remove_columns=source.column_names,
)
print(expanded.column_names)
print(len(expanded))

这里每条文本都有两个词，因此预期得到六行，列名为 word 与 parent_id。parent_id 被明确复制到每个词条，保留来源关系。remove_columns 在映射过程中移除原列，不会让回调失去读取这些输入列的机会。

输出行数变了，训练标签和样本编号也要跟着重新对应。若只扩展文本而遗漏标签，问题可能表现为长度报错，也可能成为数据含义错配，不能只让程序不报错就结束。

选择 batch_size 并识别最后一批

batch_size 控制一次交给映射函数多少条输入。上面的三条记录配 batch_size=2，且不丢弃尾批时，最后一次只剩一条；回调应根据实际列表长度处理，不要写死循环两次。默认批大小与版本参数可查 Dataset.map 参考文档，实际项目建议显式指定便于控制资源。

处理大量长文本时，先从可承受的小批开始，观察峰值内存和处理耗时再调整。更多输入行不代表同样字节量；拆分后的一批输出还可能远大于输入。num_proc 管进程并行度，writer_batch_size 管写入缓存的批量，与函数入参 batch_size 是不同层面的设置。

常见问题

为什么 batch["text"].strip() 报错？

默认批量输入的这一项是列表，不是字符串。对列表里的每个文本分别处理，或交给支持列表输入的工具。

打开 batched 后返回一个数字可以吗？

通常的列更新应返回批字典，每列是与输出批行数匹配的序列。若要给每条样本标记长度，应返回每条文本各自长度组成的列表，而不是整批的一个总数。

input_columns 会改变这个输入约定吗？

会。显式指定 input_columns 时，所选列会按位置传入回调，而非默认整批字典；对照函数签名修改，不能把两种模式混用。

最稳妥的迁移方式是先保持输入输出行数一致，再尝试拆分或扩展。看清回调接收的是一行还是列列表，并让所有输出列保持同批行数对齐，就能避免大部分 batched 入参误用。

参考：Datasets 批量映射说明、Dataset.map 参数参考。