Pyspark：在UDF中传递多个列

Question

我正在编写一个用户定义函数，它将获取除数据帧中第一个列之外的所有列并进行求和（或任何其他操作）。现在数据框有时可以有3列或4列或更多列。它会有所不同。

我知道我可以硬编码4个列名作为UDF中的传递，但在这种情况下它会有所不同所以我想知道如何完成它？

以下是第一个示例中的两个示例，我们有两列要添加，第二个示例中我们有三列要添加。

Answer 1

如果要传递给UDF的所有列都具有相同的数据类型，则可以使用array作为输入参数，例如：

>>> from pyspark.sql.types import IntegerType
>>> from pyspark.sql.functions import udf, array
>>> sum_cols = udf(lambda arr: sum(arr), IntegerType())
>>> spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B']) \
...     .withColumn('Result', sum_cols(array('A', 'B'))).show()
+---+---+---+------+
| ID|  A|  B|Result|
+---+---+---+------+
|101|  1| 16|    17|
+---+---+---+------+

>>> spark.createDataFrame([(101, 1, 16, 8)], ['ID', 'A', 'B', 'C'])\
...     .withColumn('Result', sum_cols(array('A', 'B', 'C'))).show()
+---+---+---+---+------+
| ID|  A|  B|  C|Result|
+---+---+---+---+------+
|101|  1| 16|  8|    25|
+---+---+---+---+------+

Answer 2

使用struct而不是array

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, struct
sum_cols = udf(lambda x: x[0]+x[1], IntegerType())
a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B'])
a.show()
a.withColumn('Result', sum_cols(struct('A', 'B'))).show()

Answer 3

没有数组和结构的另一种简单方法。

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, struct

def sum(x, y):
    return x + y

sum_cols = udf(sum, IntegerType())

a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B'])
a.show()
a.withColumn('Result', sum_cols('A', 'B')).show()

Answer 4

如果您不想键入所有列名，而只是将所有列转储到UDF中，则需要将列表理解包装在结构中。

from pyspark.sql.functions import struct, udf
sum_udf = udf(lambda x: sum(x[1:]))
df_sum = df.withColumn("result", sum_udf(struct([df[col] for col in df.columns])))

Answer 5

也许这是一个较晚的答案，但是我不喜欢不必要地使用UDF，所以：

from pyspark.sql.functions import col
from functools import reduce
data = [["a",1,2,5],["b",2,3,7],["c",3,4,8]]
df = spark.createDataFrame(data,["id","v1","v2",'v3'])

calculate = reduce(lambda a, x: a+x, map(col, ["v1","v2",'v3']))

df.withColumn("Result", calculate)
#
#id v1  v2  v3  Result
#a  1   2   5   8
#b  2   3   7   12
#c  3   4   8   15

在这里，您可以使用在Column中实现的任何操作。同样，如果您想编写具有特定逻辑的自定义udf，则可以使用它，因为Column提供了树执行操作。无需收集数组并求和。

如果与将过程作为数组操作进行比较，那么从性能的角度来看这将是不好的，让我们来看一下我的案例和数组案例，array案例的物理计划。

我的情况：

== Physical Plan ==
*(1) Project [id#355, v1#356L, v2#357L, v3#358L, ((v1#356L + v2#357L) + v3#358L) AS Result#363L]
+- *(1) Scan ExistingRDD[id#355,v1#356L,v2#357L,v3#358L]

数组大小写：

== Physical Plan ==
*(2) Project [id#339, v1#340L, v2#341L, v3#342L, pythonUDF0#354 AS Result#348]
+- BatchEvalPython [<lambda>(array(v1#340L, v2#341L, v3#342L))], [pythonUDF0#354]
   +- *(1) Scan ExistingRDD[id#339,v1#340L,v2#341L,v3#342L]

在可能的情况下-我们需要避免使用UDF，因为Catalyst不知道如何优化它们

Answer 6

这是我尝试并似乎起作用的方式：

colsToSum = df.columns[1:]
df_sum = df.withColumn("rowSum", sum([df[col] for col in colsToSum]))