Spark DataFrame:计算行方式(或任何聚合操作)

Question

我在内存中加载了一个Spark DataFrame,我想对列进行均值(或任何聚合操作).我该怎么办？(在numpy,这被称为进行操作axis=1).

如果有人在行(axis=0)中计算DataFrame的平均值,那么这已经内置:

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|
+------+

Answer 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表达式.例如,通过处理null来查看Spark Scala的行平均值.