apache-spark - udf多个参数 - spark withcolumn用法



Pyspark:在UDF中传递多个列 (1)

如果要传递给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|
+---+---+---+---+------+

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

我知道我可以硬编码4列名作为在UDF通过,但在这种情况下,它会有所不同,所以我想知道如何做到这一点?

在第一个例子中,我们有两列要添加,第二个例子中有三列要添加。