如何用PySpark/Pandas按列名模式匹配对多列求和生成新列

Pandas 实现 #

直接用DataFrame的filter方法配合正则匹配目标列，再沿行方向求和即可：

import pandas as pd

# 示例df构造可根据实际业务替换
# df = pd.read_csv("your_data.csv")

# 匹配规则列并逐行求和生成新列
df["plain_sum"] = df.filter(regex=r"^plain-prod.*", axis=1).sum(axis=1)

*说明：默认求和会自动忽略空值，如果需要空值参与计算时结果返回NaN，可给sum方法添加参数skipna=False。

PySpark 实现 #

先通过正则筛选出符合规则的列名，再对多列做逐行相加生成新列：

import re
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, coalesce

# 初始化SparkSession，实际场景可直接复用已有实例
spark = SparkSession.builder.appName("col_sum_calc").getOrCreate()
# 示例df构造可根据实际业务替换
# df = spark.read.csv("your_data.csv", header=True, inferSchema=True)

# 正则匹配目标列
matched_cols = [c for c in df.columns if re.match(r"^plain-prod.*", c)]
# 逐行求和，coalesce用于把空值转为0，避免空值导致求和结果为NULL
df = df.withColumn(
    "plain_sum",
    sum(coalesce(col(c), 0) for c in matched_cols)
)

*说明：如果不需要忽略空值，去掉coalesce处理逻辑，直接写sum(col(c) for c in matched_cols)即可。

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