Spark DataFrame：计算行均值（或任何聚合操作）

Posted 2023-04-15

【中文标题】Spark DataFrame：计算行均值（或任何聚合操作）

【英文标题】：Spark DataFrame: Computing row-wise mean (or any aggregate operation)

【发布时间】：2015-09-19 17:05:56

【问题描述】：

我在内存中加载了一个 Spark DataFrame，我想对列进行均值（或任何聚合操作）。我该怎么做？ （在numpy 中，这称为对axis=1 进行操作）。

如果有人正在计算 DataFrame 沿行 (axis=0) 的平均值，那么它已经内置：

from pyspark.sql import functions as F
F.mean(...)

但是有没有办法以编程方式对列中的条目执行此操作？比如下面的DataFrame

+--+--+---+---+
|id|US| UK|Can|
+--+--+---+---+
| 1|50|  0|  0|
| 1| 0|100|  0|
| 1| 0|  0|125|
| 2|75|  0|  0|
+--+--+---+---+

省略id，意思是

+------+
|  mean|
+------+
| 16.66|
| 33.33|
| 41.67|
| 25.00|
+------+

【参考方案1】：

您只需要这样的标准 SQL：

SELECT (US + UK + CAN) / 3 AS mean FROM df

可以直接用SqlContext.sql或者用DSL表达

df.select(((col("UK") + col("US") + col("CAN")) / lit(3)).alias("mean"))

如果你有更多的列，你可以生成如下表达式：

from functools import reduce
from operator import add
from pyspark.sql.functions import col, lit

n = lit(len(df.columns) - 1.0)
rowMean  = (reduce(add, (col(x) for x in df.columns[1:])) / n).alias("mean")

df.select(rowMean)

或

rowMean  = (sum(col(x) for x in df.columns[1:]) / n).alias("mean")
df.select(rowMean)

在 Scala 中终于等价了：

df.select(df.columns
  .drop(1)
  .map(col)
  .reduce(_ + _)
  .divide(df.columns.size - 1)
  .alias("mean"))

在更复杂的场景中，您可以使用 array 函数组合列并使用 UDF 来计算统计信息：

import numpy as np
from pyspark.sql.functions import array, udf
from pyspark.sql.types import FloatType

combined = array(*(col(x) for x in df.columns[1:]))
median_udf = udf(lambda xs: float(np.median(xs)), FloatType())

df.select(median_udf(combined).alias("median"))

使用 Scala API 表达的相同操作：

val combined = array(df.columns.drop(1).map(col).map(_.cast(DoubleType)): _*)
val median_udf = udf((xs: Seq[Double]) => 
    breeze.stats.DescriptiveStats.percentile(xs, 0.5))

df.select(median_udf(combined).alias("median"))

从 Spark 2.4 开始，另一种方法是将值组合到一个数组中并应用 aggregate 表达式。参见例如Spark Scala row-wise average by handling null。

【讨论】：

在你开始的例子中reduce 来自哪里？

from functools import reduce

【参考方案2】：

在 Scala 中这样的事情会做到这一点

val cols = Seq("US","UK","Can")
f.map(r => (r.getAs[Int]("id"),r.getValuesMap(cols).values.fold(0.0)(_+_)/cols.length)).toDF

【讨论】：

那不是只有在您最初没有 DataFrame 的情况下才有效吗？

no r 将获取数据框行，并且 getValueMap 提取具有 k = 列的 Map； value = 列的值，然后 values 返回一个 iterable，其中包含折叠对它们求和的值。整个事情返回一个元组 int,float 可以转回数据帧