PySpark如何实现SQL中group by col1 + cube(col2,col3)的聚合操作？

你遇到的问题是因为PySpark的GroupedData对象（groupBy返回的结果）没有cube方法，而直接对整个DataFrame调用cube会把所有传入的列都纳入多维聚合，不符合你只想让col2和col3做Cube的需求。下面提供两种可行的解决方法：

方法一：过滤Cube结果中不需要的行 #

直接调用cube(col1, col2, col3)，然后过滤掉col1为null的行——因为cube(col1, col2, col3)会生成包含col1为null的全局聚合结果，而你需要的是按col1分组后再对col2、col3做Cube，这部分结果里col1不会为null。

示例代码：

from pyspark.sql import functions as F

# 读取表
df = spark.table("table")
# 执行Cube聚合并过滤
result_df = df.cube("col1", "col2", "col3") \
              .agg(F.sum("col4").alias("sum_col4")) \
              .filter(F.col("col1").isNotNull())

这个结果和你给出的SQL语句执行效果完全一致。

方法二：手动生成Cube组合并合并（可选） #

如果你不想用过滤的方式，也可以先按col1分组，再手动生成col2和col3的所有组合聚合结果后合并。这种写法代码量更大，但适合需要更细粒度控制的场景：

from pyspark.sql import functions as F

df = spark.table("table")
# 基础分组聚合（col2、col3都不为null）
base = df.groupBy("col1", "col2", "col3").agg(F.sum("col4").alias("sum_col4"))
# 仅按col1、col2聚合（col3为null）
col2_only = df.groupBy("col1", "col2").agg(F.sum("col4").alias("sum_col4")).withColumn("col3", F.lit(None))
# 仅按col1、col3聚合（col2为null）
col3_only = df.groupBy("col1", "col3").agg(F.sum("col4").alias("sum_col4")).withColumn("col2", F.lit(None))
# 仅按col1聚合（col2、col3都为null）
col1_only = df.groupBy("col1").agg(F.sum("col4").alias("sum_col4")).withColumns({"col2": F.lit(None), "col3": F.lit(None)})
# 合并所有结果
result_df = base.unionByName(col2_only).unionByName(col3_only).unionByName(col1_only)

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