首页 > 解决方案 > PySpark 传递列表到用户定义函数

问题描述

我有一个带有一列的 DataFrame。在该列的每一行中,都有一个具有不同整数个数的列表。例如,第 1 行有一个包含 5 个整数的列表。第 2 行有一个包含 8 个整数的列表,并且......我想编写一个 UDF,它将每一行作为一个列表并将列表中的所有整数相乘并将结果作为一个数字返回。我想在另一列中有答案,所以它是 df.withColumn(...) 出于某种原因,我确实想使用 UDF而不是任何其他预构建的函数或命令。感谢您的时间和支持。

标签: listdataframepysparkbigdatauser-defined-functions

解决方案


定义一些示例输入数据:

df = spark.createDataFrame([
    (1, [3, 4, 8]), (2, [7, 2, 6, 8])
], ("id", "list_of_ints"))

df.show()

的定义udf

from pyspark.sql.functions import udf, col
def product(numbers):
     a = 1
     for num in numbers:
         a *= num
     return a

from pyspark.sql.types import IntegerType
product_udf = udf(lambda z: product(z), IntegerType())

并添加一个包含所有列表元素乘积的列:

df.withColumn("product", product_udf("list_of_ints")).show()

+---+------------+-------+
| id|list_of_ints|product|
+---+------------+-------+
|  1|   [3, 4, 8]|     96|
|  2|[7, 2, 6, 8]|    672|
+---+------------+-------+

希望这可以帮助!


推荐阅读