如何对涉及Scipy稀疏矩阵的Pandas语句进行向量化优化

使用Pandas编写数据处理逻辑时，apply方法的运行耗时最长可达2分钟。根据性能优化相关资料，将逻辑改写为向量化形式可大幅提升运行速度。
原始实现代码如下：

output_data["on_s"] = output_data["m_ind"].apply(lambda x: my_matrix[x, 0] + my_matrix[x, 1] + my_matrix[x, 2])

其中my_matrix为scipy.sparse类型稀疏矩阵。最初的优化思路是调用sum方法对矩阵按行求和：

summed_matrix = my_matrix.sum(axis=1)

但完成该步骤后，无法衔接后续逻辑匹配目标数据。

补充样例数据 #

• my_matrix为scipy.sparse.csr_matrix格式，样例数据如下：

(290730, 2)     0.3058016922838267
(290731, 2)     0.3390328430763723
(290733, 2)     0.0838999800585995
(290734, 2)     0.0237008960604337
(290735, 2)     0.0116864263235209

• output_data["m_ind"]为Pandas Series类型，样例取值如下：

97543
97544
97545
97546
97547

你的按行求和思路是正确的，核心问题是没有利用稀疏矩阵的批量索引能力，避免逐行Python循环。
需求本质是：提取my_matrix中每行第0、1、2列的和，再根据m_ind存储的行号，批量匹配对应结果赋值给新列，全程不需要apply循环。
具体实现代码如下：

# 第一步：对稀疏矩阵所有行的0、1、2列做行向求和，得到形状为(矩阵总行数, 1)的结果
# 如果矩阵总列数恰好为3，可直接简写为 my_matrix.sum(axis=1)
summed_result = my_matrix[:, [0, 1, 2]].sum(axis=1)
# 第二步：将求和结果转为一维numpy数组，用m_ind的取值批量索引，直接赋值给目标列
# .A1是scipy稀疏矩阵/矩阵对象的属性，用于快速转为一维numpy数组，避免维度不匹配报错
output_data["on_s"] = summed_result.A1[output_data["m_ind"].values]

性能说明 #

• 原apply实现为Python层逐行循环，每次循环都要单独执行稀疏矩阵索引操作，大量耗时消耗在循环调度和单次索引的 overhead 上，这也是耗时达到2分钟的核心原因
• 上述实现全程为底层C级别的向量化运算，无Python层显式循环，通常可以将耗时压缩到毫秒级。

注意：如果m_ind中存在超出my_matrix总行数的索引值，需要先做索引合法性校验，避免取值报错。

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