时间戳

首页 > 解决方案 > 将函数应用于 RDD 中的列（python、spark）

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

问题描述

这是我的RDD：

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

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

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

标签： pythonlistapache-sparkpysparkrdd

解决方案

为了将其应用于数据框上的列，您可以创建并应用用户定义函数（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")
)

推荐阅读