PySpark传递给用户定义功能的列表

Question

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

Answer 1

定义一些样本输入数据：

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

希望这会有所帮助！