pySpark groupby 中的条件聚合

Question

pySpark 新手的简单问题：我有一个 df，我想进行条件聚合，如果分母不同于 0，则返回聚合结果，否则为 0。

我的尝试产生错误：

groupBy=["K"]
exprs=[(sum("A")+(sum("B"))/sum("C") if sum("C")!=0 else 0 ]
grouped_df=new_df.groupby(*groupBy).agg(*exprs)

有什么提示吗？

谢谢

Answer 1

您必须为 if/else 使用 when/otherwise：

import pyspark.sql.functions as psf
new_df.groupby("K").agg(
    psf.when(psf.sum("C")==0, psf.lit(0)).otherwise((psf.sum("A") + psf.sum("B"))/psf.sum("C")).alias("sum")
)

不过你也可以这样做：

import pyspark.sql.functions as psf
new_df.groupby("K").agg(
    ((psf.sum("A") + psf.sum("B"))/psf.sum("C")).alias("sum")
).na.fill({"sum": 0})

pySpark groupby 中的条件聚合

conditional aggragation in pySpark groupby

pyspark-sql