在 PySpark 的 DataFrame 中聚合列数组? [英] Aggregate over column arrays in DataFrame in PySpark?

问题描述

假设我有以下 DataFrame:

[Row(user='bob', values=[0.5, 0.3, 0.2]),
Row(user='bob', values=[0.1, 0.3, 0.6]),
Row(user='bob', values=[0.8, 0.1, 0.1])]

我想 groupBy user 并执行类似 avg(values) 的操作，其中对数组的每个索引取平均值 <代码>值像这样:

I would like to groupBy user and do something like avg(values) where the average is taken over each index of the array values like this:

[Row(user='bob', avgerages=[0.466667, 0.233333, 0.3])]

如何在 PySpark 中执行此操作?

How can I do this in PySpark?

推荐答案

您可以扩展数组并计算每个索引的平均值.

You can expand array and compute average for each index.

Python

from pyspark.sql.functions import array, avg, col

n = len(df.select("values").first()[0])

df.groupBy("user").agg(
    array(*[avg(col("values")[i]) for i in range(n)]).alias("averages")
)

Scala

import spark.implicits._
import org.apache.spark.functions.{avg, size}

val df = Seq(
  ("bob", Seq(0.5, 0.3, 0.2)),
  ("bob", Seq(0.1, 0.3, 0.6))
).toDF("user", "values")

val n = df.select(size($"values")).as[Int].first
val values = (0 to n).map(i => $"values"(i))

df.select($"user" +: values: _*).groupBy($"user").avg()

这篇关于在 PySpark 的 DataFrame 中聚合列数组?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！