使用dask.bag.map时arg重复分配内存的解决办法咨询

你遇到的问题根源是：使用bag.map传递外部大对象arg时，Dask会为每个元素级任务序列化并复制一份arg，当分区/元素数量多、arg体积大时，必然导致内存占用飙升。下面是几种实用的解决方法：

1. 改用map_partitions批量处理分区 #

map_partitions是针对整个分区做批量操作，每个分区仅需传入一次arg，复制次数从「元素数量」降到「分区数量」，能大幅减少内存开销。

修改后的代码示例：

from pathlib import Path
import numpy as np
import pandas as pd
from dask import bag

# 目录与数据生成代码保持不变
in_dir = Path.home() / 'in_dir'
out_dir = Path.home() / 'out_dir'
in_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents=True, exist_ok=True)

n_files = 100
n_lines_per_file = int(1e6)
df = pd.DataFrame({'a': np.arange(n_lines_per_file).astype(str)})
for i in range(n_files):
    df.to_csv(in_dir / f'{i}.txt', index=False, header=False)

def mapper(x, arg):
    y = x  # 替换为你的实际处理逻辑
    return y

def partition_mapper(partition, arg):
    # 对整个分区的所有元素批量应用mapper
    return (mapper(x, arg) for x in partition)

arg = np.zeros(int(1e7))

(
    bag
    .read_text(str(in_dir / '*.txt'))
    .map_partitions(partition_mapper, arg)
    .to_textfiles(str(out_dir / '*.txt'))
)

2. 分布式集群下用广播共享对象 #

如果使用Dask Distributed集群，可通过client.scatter将arg广播到每个工作节点，每个节点仅保留一份副本，所有任务共享该内存对象：

from dask.distributed import Client
from pathlib import Path
import numpy as np
import pandas as pd
from dask import bag

# 初始化分布式客户端
client = Client()

# 目录与数据生成代码不变
in_dir = Path.home() / 'in_dir'
out_dir = Path.home() / 'out_dir'
in_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents=True, exist_ok=True)

n_files = 100
n_lines_per_file = int(1e6)
df = pd.DataFrame({'a': np.arange(n_lines_per_file).astype(str)})
for i in range(n_files):
    df.to_csv(in_dir / f'{i}.txt', index=False, header=False)

def mapper(x, arg):
    y = x  # 替换为你的实际处理逻辑
    return y

arg = np.zeros(int(1e7))
# 广播arg到所有工作节点，每个节点仅存一份
arg_future = client.scatter(arg, broadcast=True)

(
    bag
    .read_text(str(in_dir / '*.txt'))
    .map(mapper, arg_future)
    .to_textfiles(str(out_dir / '*.txt'))
).compute()

3. 单机场景下使用全局变量（快速临时方案） #

在单机多进程模式下，可将arg设为全局变量，利用操作系统的写时复制（Copy-on-Write）机制，只读场景下不会重复分配内存：

from pathlib import Path
import numpy as np
import pandas as pd
from dask import bag

# 目录与数据生成代码不变
in_dir = Path.home() / 'in_dir'
out_dir = Path.home() / 'out_dir'
in_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents=True, exist_ok=True)

n_files = 100
n_lines_per_file = int(1e6)
df = pd.DataFrame({'a': np.arange(n_lines_per_file).astype(str)})
for i in range(n_files):
    df.to_csv(in_dir / f'{i}.txt', index=False, header=False)

# 全局定义arg，仅在单机场景使用
arg = np.zeros(int(1e7))

def mapper(x):
    # 直接使用全局arg
    y = x  # 替换为你的实际处理逻辑
    return y

(
    bag
    .read_text(str(in_dir / '*.txt'))
    .map(mapper)
    .to_textfiles(str(out_dir / '*.txt'))
)

注意：此方法仅适用于单机，且arg必须保持只读，修改arg会触发写时复制，导致内存占用回升。

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