如何为Dask Bag中各元素应用不同函数（类to_textfiles逻辑）

核心思路就是把Bag里的每个元素和唯一索引配对，这样在map的时候就能根据索引生成对应的操作逻辑，具体做法分几步：

1. 生成索引序列的Bag #

先拿到原Bag的总元素数，再生成一个包含连续索引的Dask Bag，保持分区数和原Bag一致，避免后续配对时打乱分区：

import dask.bag as db

# 假设你的原Bag是original_bag
total_count = original_bag.count().compute()
index_bag = db.from_sequence(range(total_count), npartitions=original_bag.npartitions)

2. 把原Bag和索引Bag配对 #

用db.zip把两个Bag合并，每个元素会变成(索引值, 原元素)的元组：

indexed_bag = db.zip(index_bag, original_bag)

3. 在map里用索引实现自定义逻辑 #

现在就能在map操作里同时拿到索引和元素，按需生成对应的函数或路径：

示例1：按索引做乘法 #

def multiply_with_index(index, num):
    return num * index

# 对配对后的元素执行操作
result_bag = indexed_bag.map(lambda item: multiply_with_index(item[0], item[1]))

示例2：保存numpy数组到对应路径 #

import numpy as np

def save_array(index, arr):
    save_path = f"/foo/bar/{index}.npy"
    np.save(save_path, arr)
    return save_path  # 可选返回路径，方便后续验证

# 触发执行（Dask是惰性计算，必须调用compute()才会实际保存）
saved_paths = indexed_bag.map(lambda item: save_array(item[0], item[1])).compute()

适配未知长度的Bag #

如果原Bag是动态生成的，没法提前拿到总长度，可以用map_partitions给每个分区内的元素加本地索引，再加上分区的起始偏移量：

def add_offset(partition, offset):
    # 给分区内每个元素加上分区起始偏移，生成全局索引
    return [(offset + idx, elem) for idx, elem in enumerate(partition)]

# 计算每个分区的起始偏移量
offsets = []
current_offset = 0
for part in original_bag.partitions:
    offsets.append(current_offset)
    current_offset += part.count().compute()

# 给每个分区应用偏移，再扁平化结果
indexed_bag = original_bag.map_partitions(add_offset, offsets).flatten()

这种方式不用提前算总长度，适合流式或者不确定元素数量的场景。

内容的提问来源于stack exchange，提问作者hikaru_alpha