PySpark中如何使用子查询？解决分组统计报错问题

需求与原SQL逻辑 #

想要实现的SQL查询逻辑如下（注：原SQL存在笔误，ODER应为ORDER，且HAVING条件逻辑有误，正确应为筛选分组计数等于最小分组计数的记录）：

SELECT age, COUNT(age)
FROM T
GROUP BY age
HAVING COUNT(age) = (SELECT MIN(count_age) FROM (SELECT COUNT(age) AS count_age FROM T GROUP BY age) t)
ORDER BY COUNT(age)

尝试的代码及报错 #

尝试了以下PySpark代码：

min_size = df.groupBy("age").count().select(f.min("count"))
df.groupBy("age").count().sort("count").filter(f.col("count")==min_size).show()

触发报错：

AttributeError: 'DataFrame' object has no attribute '_get_object_id'

解答 #

PySpark是否支持子查询？ #

完全支持，不管是DataFrame API还是Spark SQL语法，都可以使用子查询实现复杂逻辑。

错误原因 #

报错的核心原因是min_size是一个DataFrame对象，而filter方法需要的是列表达式或单个标量值，直接将DataFrame与列进行比较会导致类型不匹配，PySpark无法解析这种操作。

解决方法 #

提供三种可行的解决方式：

方式一：提取标量值后过滤

先从统计结果中提取出最小分组计数的具体数值，再用于过滤：

# 计算最小分组计数并提取为标量
min_size = df.groupBy("age").count().select(f.min("count")).first()[0]
# 分组统计后筛选出计数等于最小值的记录
df.groupBy("age").count() \
  .filter(f.col("count") == min_size) \
  .sort("count") \
  .show()

方式二：使用窗口函数（大数据场景更高效）

通过窗口函数一次性计算全局最小分组计数，避免多次分组操作，性能更优：

from pyspark.sql import Window

# 定义全局窗口（所有数据为一个分区）
global_window = Window.partitionBy()
df.groupBy("age").count() \
  .withColumn("min_global_count", f.min("count").over(global_window)) \
  .filter(f.col("count") == f.col("min_global_count")) \
  .drop("min_global_count") \
  .sort("count") \
  .show()

方式三：使用Spark SQL语法

如果更习惯SQL写法，可以直接注册临时视图后执行SQL：

# 将DataFrame注册为临时视图
df.createOrReplaceTempView("T")
# 执行修正后的SQL查询
spark.sql("""
    SELECT age, COUNT(age) AS count_age
    FROM T
    GROUP BY age
    HAVING count_age = (
        SELECT MIN(sub_count) 
        FROM (SELECT COUNT(age) AS sub_count FROM T GROUP BY age) sub_t
    )
    ORDER BY count_age
""").show()

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