Pyspark: передать несколько столбцов в UDF

Я пишу пользовательскую функцию, которая будет принимать все столбцы, кроме первого, в фреймворке данных и делать сумму (или любую другую операцию). Теперь в dataframe иногда могут быть 3 столбца или 4 столбца или более. Он будет меняться.

Я знаю, что могу жестко кодировать 4 имени столбца в качестве прохода в 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|
+---+---+---+---+------+

Ответ 2

Использовать структуру вместо массива

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()

Ответ 3

Еще один простой способ без Array и Struct.

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()

Ответ 4

Вот как я пытался и, похоже, работал:

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