PySpark：单个.agg()调用中是否允许存在依赖关系的聚合操作？

你编写的PySpark代码示例如下：

df = ... # some dataframe
res_df = df.groupby('col1', 'col2').agg(
    <aggregation func/expr>.alias('agg1'),
    <aggregation func/expr that depends on agg1>.alias('agg2')
)

典型场景比如用collect_list()生成agg1列，再用size()基于agg1计算agg2列。

关于这种写法的跨版本兼容性，结论如下：

• Spark 3.0及以上版本：完全支持这种写法。Spark 3.0引入的优化器能够识别agg()内部的列依赖关系，自动调整执行顺序——先计算出agg1，再基于它计算agg2，所以你的代码能正常运行。
• Spark 3.0以下版本：不支持这种写法。旧版本的Spark在分析阶段无法解析agg()内部引用的agg1列，会直接抛出“无法解析列”的错误，此时必须拆分操作，先完成第一次聚合得到agg1，再通过后续的select或withColumn计算agg2。

GitHub Copilot Review给出报错提示，大概率是基于旧版本Spark的行为给出的判断，没有考虑到你当前使用的是较新版本的Spark。

如果要兼顾可读性和跨版本兼容，可按以下方式处理：

• 仅需兼容Spark 3.0+：保留当前写法，既保证可读性，又能被优化器正确处理。
• 需要兼容低版本：拆分操作，示例代码如下：

from pyspark.sql.functions import col

res_df = df.groupby('col1', 'col2').agg(
    <aggregation func/expr>.alias('agg1')
).withColumn('agg2', <func>(col('agg1')))

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