Pyspark:在UDF中传递多列 [英] Pyspark: Pass multiple columns in UDF

问题描述

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

I am writing a User Defined Function which will take all the columns except the first one in a dataframe and do sum (or any other operation). Now the dataframe can sometimes have 3 columns or 4 columns or more. It will vary.

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

I know I can hard code 4 column names as pass in the UDF but in this case it will vary so I would like to know how to get it done?

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

Here are two examples in the first one we have two columns to add and in the second one we have three columns to add.

推荐答案

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

If all columns you want to pass to UDF have the same data type you can use array as input parameter, for example:

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

这篇关于Pyspark:在UDF中传递多列的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！