pyspark - 使用 ArrayType 列折叠和求和

如何解决pyspark - 使用 ArrayType 列折叠和求和

我正在尝试按元素求和，并且我创建了这个虚拟 df。输出应该是 [10,4,1]

from pyspark.sql.types import StructType,StructField,StringType,IntegerType,ArrayType
data = [
    ("James",[1,1,1]),("James",[2,0]),[3,[4,0])
  ]

schema = StructType([ \
    StructField("firstname",StringType(),True),\
    StructField("scores",ArrayType(IntegerType()),True) \
  ])
 
df = spark.createDataFrame(data=data,schema=schema)

posexplode 有效，但我的真实 df 太大，所以我尝试使用 fold，但它给了我一个错误。有任何想法吗？谢谢！

vec_df = df.select("scores")
vec_sums = vec_df.rdd.fold([0]*4,lambda a,b: [x + y for x,y in zip(a,b)])

文件“”，第 2 行，

TypeError: 不支持 + 的操作数类型：'int' 和 'list'

解决方法

您需要在 fold 之前将行的 RDD 映射到列表的 RDD：

vec_sums = vec_df.rdd.map(lambda x: x[0]).fold([0]*4,lambda a,b: [x + y for x,y in zip(a,b)])

为了帮助理解，您可以查看 RDD 的外观。

>>> vec_df.rdd.collect()
[Row(scores=[1,1,1]),Row(scores=[2,0]),Row(scores=[3,Row(scores=[4,0])]

>>> vec_df.rdd.map(lambda x: x[0]).collect()
[[1,1],[2,0],[3,[4,0]]

因此您可以想象 vec_df.rdd 包含一个嵌套列表，需要在 fold 之前取消嵌套。