派斯帕克：在 UDF 中传递多列

Question

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

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

这里有两个示例，第一个示例中我们有两列要添加，第二个示例中有三列要添加。

Answer 1

使用结构而不是数组

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 2

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

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

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 3

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

>>> 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|
+---+---+---+---+------+