使用 pyspark/pandas 使用列名的模式匹配对具有不同列名的多个列值求和

Posted 2023-04-15

【中文标题】使用 pyspark/pandas 使用列名的模式匹配对具有不同列名的多个列值求和

【英文标题】：summing multiple columns values which has different column names by using pattern matching of column names using pyspark/pandas

【发布时间】：2021-11-25 15:26:06

【问题描述】：

例如： 我有一个数据框，其中的列像

lens	plain-prod 102	plain-prod 105	plain-prod 107
First	1	3	4
Second	2	5	3
First	3	7	2
Second	4	8	1

所以我需要进行模式匹配 (^plain-prod.*) 并选择所有匹配的 3 列并创建新列 plain_sum，其中包含如何使用 pyspark 或 pandas 实现此目的的总和。

lens	plain-prod 102	plain-prod 105	plain-prod 107	plain_sum
First	1	3	4	8
Second	2	5	3	10
First	3	7	2	12
Second	4	8	1	13

【参考方案1】：

用 Pandas 试试这个方法（df 是你的数据框）：

df['plain_sum'] = df.filter(regex='^plain-prod.*').sum(axis=1)

【参考方案2】：

你可以在 PySpark 中做这样的事情：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.getOrCreate()
data = [
    "lens": "First", "plain-prod 102": 1, "plain-prod 105": 3, "plain-prod 107": 4,
    "lens": "Second", "plain-prod 102": 2, "plain-prod 105": 3, "plain-prod 107": 4,
]
df = spark.createDataFrame(data)
df = df.withColumn(
    "plain_sum", sum([F.col(x) for x in df.columns if "plain-prod" in x])
)

结果：

+------+--------------+--------------+--------------+---------+                 
|lens  |plain-prod 102|plain-prod 105|plain-prod 107|plain_sum|
+------+--------------+--------------+--------------+---------+
|First |1             |3             |4             |8        |
|Second|2             |3             |4             |9        |
+------+--------------+--------------+--------------+---------+

【讨论】：

这会失败并返回错误 Invalid argument, not a string or column

@Ajaykumar 请提供您正在运行的代码示例。

我只是通过更改数据框名称和过滤器运行相同的代码

@Ajaykumar 我猜您的原始数据集与我的测试示例不同。无论如何，我已经更新了答案，如果您需要它，也许会对您有所帮助。

你在 pyspark 中执行了你的代码吗？如果不尝试，您将收到您创建的示例数据框的错误消息