Spark如何在where条件中使用聚合函数？

为什么原代码无法运行 #

f.avg("age") 是聚合函数，而 where 子句是对单条原始行进行过滤，此时还未执行任何聚合计算，因此不能直接在 where 中使用聚合函数。必须先计算出平均年龄，再将其作为过滤条件使用。

正确的PySpark实现方式 #

方法1：先计算全局平均年龄，再过滤

这种方式和你写的Spark SQL逻辑最接近：

from pyspark.sql import functions as f

# 计算整个数据集的平均年龄
global_avg_age = titan.select(f.avg("age")).first()[0]

# 执行过滤与聚合
(
    titan
    .where("Survived = 1")
    .where(f.col("age") > global_avg_age)
    .groupBy('Sex', 'Pclass')
    .count()
    .show()
)

方法2：使用窗口函数计算全局平均

通过窗口函数将全局平均年龄添加为列，再进行过滤：

from pyspark.sql import Window

# 创建覆盖所有行的窗口
all_rows_window = Window.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

(
    titan
    .where("Survived = 1")
    .withColumn("global_avg_age", f.avg("age").over(all_rows_window))
    .where(f.col("age") > f.col("global_avg_age"))
    .groupBy('Sex', 'Pclass')
    .count()
    .show()
)

方法3：在过滤条件中直接嵌套子查询

完全复刻Spark SQL的子查询写法：

(
    titan
    .where("Survived = 1")
    .where(f.col("age") > titan.select(f.avg("age")).alias("avg_age"))
    .groupBy('Sex', 'Pclass')
    .count()
    .show()
)

以上三种方法都能实现你Spark SQL代码的功能。

内容的提问来源于stack exchange，提问作者Astora