将列总和添加为PySpark数据框中的新列

Question

我正在使用PySpark，并且有一个带有一堆数字列的Spark数据框。我想添加一列，它是所有其他列的总和。

假设我的数据框具有列“ a”，“ b”和“ c”。我知道我可以这样做：

df.withColumn('total_col', df.a + df.b + df.c)

问题是我不想单独键入每列并添加它们，尤其是当我有很多列时。我希望能够自动执行此操作，或者希望通过指定要添加的列名列表来执行此操作。还有另一种方法吗？

Answer 1

这并不明显。我看不到spark Dataframes API中定义的列的基于行的总和。

这可以用相当简单的方式完成：

newdf = df.withColumn('total', sum(df[col] for col in df.columns))

df.columns由pyspark作为字符串列表提供，这些字符串给出了Spark
Dataframe中的所有列名称。对于不同的总和，您可以提供其他任何列名称列表。

我没有尝试将其作为第一个解决方案，因为我不确定它的行为。但这有效。

版本1

这太复杂了，但是也可以。

你可以这样做：

1. 用于df.columns获取列名称的列表
2. 使用该名称列表创建列的列表
3. 将该列表传递给将以折叠类型的函数方式调用列的重载add函数的东西

使用python的reduce，有关运算符重载如何工作的一些知识，以及此处列的pyspark代码变为：

def column_add(a,b):
     return  a.__add__(b)

newdf = df.withColumn('total_col', 
         reduce(column_add, ( df[col] for col in df.columns ) ))

请注意，这是python减少，而不是spark RDD减少，第二个参数reduce中的括号项需要括号，因为它是一个列表生成器表达式。

经过测试，有效！

$ pyspark
>>> df = sc.parallelize([{'a': 1, 'b':2, 'c':3}, {'a':8, 'b':5, 'c':6}, {'a':3, 'b':1, 'c':0}]).toDF().cache()
>>> df
DataFrame[a: bigint, b: bigint, c: bigint]
>>> df.columns
['a', 'b', 'c']
>>> def column_add(a,b):
...     return a.__add__(b)
...
>>> df.withColumn('total', reduce(column_add, ( df[col] for col in df.columns ) )).collect()
[Row(a=1, b=2, c=3, total=6), Row(a=8, b=5, c=6, total=19), Row(a=3, b=1, c=0, total=4)]