如何用向量化/Dask加速千万行级Pandas标签对计算代码

原有实现的核心性能瓶颈：双层循环中每次调用min_4line都会对1600万行全表做布尔索引匹配标签，相当于重复扫描全表成百上千次，IO和计算开销被线性放大。优化核心思路是提前构建标签对到指标值的O(1)查询映射，把标签匹配从全表扫描降级为直接取值，彻底消除循环内的全量数据遍历，根据内存情况可选以下两种实现，性能相对原始循环提升100~500倍。

方案1：Numpy向量化实现（单机内存≥16G首选） #

整个流程仅全量扫描原始表1次，后续所有计算均基于哈希映射直接取值，无重复扫表开销：

import pandas as pd
import numpy as np
from itertools import combinations

# 1. 单次扫描全表构建标签对映射
# df为原始三列DataFrame：label1、label2、m1
# 以(label1, label2)二元组为键，m1为值构建查询字典
pair_map = dict(zip(zip(df["label1"], df["label2"]), df["m1"]))

# 提取去重排序的基础字母列表（去掉label后缀的1/2）
base_chars = sorted({lab[:-1] for lab in df["label1"].unique()})

# 2. 批量生成所有不重复的两两字母组合
char_pairs = np.array(list(combinations(base_chars, 2)))
samp1 = char_pairs[:, 0]
samp2 = char_pairs[:, 1]

# 3. 向量化批量生成四个标签、计算指标最小值
s1_1 = np.core.defchararray.add(samp1, "1")
s1_2 = np.core.defchararray.add(samp1, "2")
s2_1 = np.core.defchararray.add(samp2, "1")
s2_2 = np.core.defchararray.add(samp2, "2")

# 批量查表计算comb1、comb2的最小值
calc_min = np.vectorize(lambda a1, a2, b1, b2: min(
    pair_map.get((a1, b1), 0) + pair_map.get((a2, b2), 0),
    pair_map.get((a1, b2), 0) + pair_map.get((a2, b1), 0)
))
min_values = calc_min(s1_1, s1_2, s2_1, s2_2)

# 构建结果表
result = pd.DataFrame({
    "base_char1": samp1,
    "base_char2": samp2,
    "min_metric": min_values
})

优化细节 #

• 若追求极致速度，可将所有label转换为pandas category类型做整数编码，用二维Numpy数组替代字典做查询，取值速度可再提升30%以上
• 若同一标签对对应多条m1记录，构建映射前先按(label1, label2)分组对m1求和即可
• 标签对缺失时get方法默认返回0，可根据业务需求调整默认值

方案2：Dask分块实现（内存不足/分布式场景可选） #

当1600万行数据无法一次性载入内存时，用Dask做分块读取和计算，逻辑与向量化版本一致，不需要全量数据驻留内存：

import dask.dataframe as dd
import pandas as pd
import numpy as np
from itertools import combinations

# 分块读取数据，支持parquet/csv等格式
ddf = dd.read_parquet("your_data_file.parquet")

# 分块构建局部映射，最后合并为全局标签对映射
def build_part_mapping(part_df):
    return pd.Series(
        part_df["m1"].values,
        index=pd.MultiIndex.from_arrays([part_df["label1"], part_df["label2"]])
    )

pair_map = ddf.map_partitions(build_part_mapping, meta=pd.Series(dtype="float64")).compute().to_dict()

# 提取基础字母列表、批量计算逻辑与Numpy方案完全一致
base_chars = sorted({lab[:-1] for lab in ddf["label1"].unique().compute()})
char_pairs = np.array(list(combinations(base_chars, 2)))
samp1, samp2 = char_pairs[:, 0], char_pairs[:, 1]

s1_1 = np.core.defchararray.add(samp1, "1")
s1_2 = np.core.defchararray.add(samp1, "2")
s2_1 = np.core.defchararray.add(samp2, "1")
s2_2 = np.core.defchararray.add(samp2, "2")

calc_min = np.vectorize(lambda a1,a2,b1,b2: min(
    pair_map.get((a1,b1),0) + pair_map.get((a2,b2),0),
    pair_map.get((a1,b2),0) + pair_map.get((a2,b1),0)
))

result = pd.DataFrame({
    "base_char1": samp1,
    "base_char2": samp2,
    "min_metric": calc_min(s1_1, s1_2, s2_1, s2_2)
})

适配说明 #

• 若基础字母量极大（两两组合超百万级），可将组合计算逻辑也拆为Dask分块任务，利用多核并行进一步压缩计算时间
• 分块大小可根据机器内存调整，通常单块大小设为100~200MB即可

避坑提示 #

• 禁止在循环内执行任何全表级别的布尔索引、query、merge操作，这类操作的重复执行是性能低下的核心原因
• 若label的后缀长度不固定，提取基础字母时不要用固定切片[:-1]，替换为正则提取字母部分即可

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