PySpark数据帧上分组数据的Pandasstyle变换

Question

如果我们有一个由一列类别和一列值组成的Pandas数据框，我们可以通过执行以下操作来删除每个类别中的平均值：

df["DemeanedValues"] = df.groupby("Category")["Values"].transform(lambda g: g - numpy.mean(g))

据我所知，Spark dataframes不直接提供这种分组/转换操作（我在spark1.5.0上使用PySpark）。那么，实现这种计算的最佳方法是什么？在

我尝试使用groupby/join，如下所示：

^{pr2}$

但这是非常缓慢的，因为，据我所知，每个类别都需要对数据帧进行完整扫描。在

我认为（但尚未验证）如果我将分组结果按/平均值收集到字典中，然后在UDF中使用该字典，则可以大大加快速度，如下所示：

nameToMean = {...}
f = lambda category, value: value - nameToMean[category]
categoryDemeaned = pyspark.sql.functions.udf(f, pyspark.sql.types.DoubleType())
df = df.withColumn("DemeanedValue", categoryDemeaned(df.Category, df.Value))

有没有一种惯用的方法来表达这种类型的操作而不牺牲性能呢？在

Answer 1

您可以使用Window来执行此操作

即

import pyspark.sql.functions as F
from pyspark.sql.window import Window

window_var = Window().partitionBy('Categroy')
df = df.withColumn('DemeanedValues', F.col('Values') - F.mean('Values').over(window_var))

Answer 2

实际上，在Spark中有一种惯用的方法来实现这一点，使用HiveOVER表达式。在

即

df.registerTempTable('df')
with_category_means = sqlContext.sql('select *, mean(Values) OVER (PARTITION BY Category) as category_mean from df')

在发动机罩下，这是使用车窗功能。不过，我不确定这是否比你的解决方案快

Answer 3

I understand, each category requires a full scan of the DataFrame.

不，它不是。数据帧聚合是使用类似于aggregateByKey的逻辑执行的。请参阅DataFrame groupBy behaviour/optimization较慢的部分是join，它需要排序/洗牌。但它仍然不需要每个组扫描。在

如果这是一个确切的代码，那么使用它是很慢的，因为您没有提供连接表达式。因为它只是执行笛卡尔积。因此，它不仅效率低下，而且是不正确的。你想要这样的东西：

from pyspark.sql.functions import col

means = df.groupBy("Category").mean("Values").alias("means")
df.alias("df").join(means, col("df.Category") == col("means.Category"))

I think (but have not verified) that I can speed this up a great deal if I collect the result of the group-by/mean into a dictionary, and then use that dictionary in a UDF

这是可能的，尽管性能会因具体情况而有所不同。使用Python udf的一个问题是它必须在Python之间来回移动数据。不过，这绝对值得一试。但是，您应该考虑为nameToMean使用广播变量。在

Is there an idiomatic way to express this type of operation without sacrificing performance?

在PySpark 1.6中，您可以使用broadcast函数：

^{pr2}$

但在<；=1.5中不可用。在