在PySpark数据框中添加列和作为新列

Question

我用的是PySpark，我有一个Spark数据框架，里面有一堆数字列。我想添加一列，它是所有其他列的总和。

假设我的数据帧具有列“a”、“b”和“c”。我知道我能做到：

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

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

Answer 1

解决方案

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

由@Paul works发布。尽管如此，我还是犯了错误，正如我所看到的那样，

TypeError: 'Column' object is not callable

一段时间后，我发现了问题（至少在我的情况下）。问题是，我之前导入了一些pyspark函数

from pyspark.sql.functions import udf, col, count, sum, when, avg, mean, min

因此，该行导入了总和 pyspark 命令，而 df.withColumn（“总计”，总和（df[col] 表示 df.列中的 col））应该使用正常的 python 和函数。

您可以使用＜code＞del sum＜code＞删除pyspark函数的引用。

否则，在我的例子中，我将导入更改为

import pyspark.sql.functions as F

然后将函数引用为F.sum。

Answer 2

最直接的方法是使用 expr 函数

from pyspark.sql.functions import *
data = data.withColumn('total', expr("col1 + col2 + col3 + col4"))

Answer 3

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

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

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

< code>df.columns由pyspark提供，作为一个字符串列表，给出spark数据帧中所有列的名称。对于不同的sum，您可以提供任何其他列名列表。

我没有尝试将此作为我的第一个解决方案，因为我不确定它将如何运行。但它是有效的。

这过于复杂，但也很有效。

您可以这样做:

1. 使用 df.列获取列名称的列表
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化简，化简的第二个参数中的括号项需要括号，因为它是一个列表生成器表达式。

经过测试，有效！

$ 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)]