将函数应用于 RDD 中的列python、spark

编程问答 2022-05-15

问题描述

这是我的RDD：

id|               arr |
+--------------------+-
|1|  [8,5,1,11,10,8,2]|
|2|    [3,6,3,2]|
|3|    [4,2,3]|
|4|    [0,0]|
|5|    [3,4,7,2]|
|6|    [1,0]|
|7|    [2,9,0]|
|8|    [3,3]|
|9| [1,12,5]|

我正在研究如何应用一个函数来对列表中的所有数字求和并在单独的列中返回总和。这是我的功能（我使用python）。它适用于一个数组，但我不知道如何将它应用于 RDD 中的一列。

def sum_func(x):
  t = 0
  for i in range(0,len(x)):
    t = t + x[i]
  return t == 0

解决方法

为了将其应用于数据帧上的列，您可以创建并应用用户定义函数 (UDF)。

from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

def sum_func(x):
  t = 0
  for i in range(0,len(x)):
    t = t + x[i]
  return t

# Creating the UDF with return type Integer

sum_func_udf = udf(sum_func,IntegerType())

然后在您的数据帧上（假设它存储在 df 中），我们使用 withColumn 添加另一列

df = df.withColumn(
   sum_func_udf(df.arr).alias("sum")
)
apache-sparklistlistpysparkpysparkpythonrdd