在PySpark DataFrame中按时间窗口计算几何均值

核心思路 #

PySpark没有内置的几何均值函数，但可以通过对数转换+算术平均+指数还原的方式实现：
几何均值 = exp( 平均值( log(Val1) ) )
这个逻辑等价于n个数值乘积的n次方根，完全符合几何均值的数学定义。

实现代码 #

方式1：分组聚合得到目标结果（匹配你需要的输出格式）

from pyspark.sql import functions as F

# 按Time分组计算几何均值
result_df = df.groupBy("Time") \
    .agg(
        F.exp(F.avg(F.log(F.col("Val1")))).alias("geo_mean")
    ) \
    .orderBy("Time")

# 查看结果
result_df.show(truncate=False)

方式2：用窗口函数添加几何均值列（和你计算算术平均的方式一致）

如果需要保留原DataFrame的所有行，同时添加对应Time分组的几何均值：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 定义窗口：按Time分组
w = Window.partitionBy("Time")

# 添加几何均值列
df_with_geo = df.withColumn(
    "geo_mean",
    F.exp(F.avg(F.log(F.col("Val1"))).over(w))
)

df_with_geo.show(truncate=False)

结果验证 #

• Time=1时：3×6×8=144，三次方根≈5.24148，与预期结果一致
• Time=2时：10×21×33=6930，三次方根≈19.06533，与预期结果一致

注意事项 #

如果数据中存在0或负数，log函数会报错，需要提前过滤这类数据，或者根据业务场景做特殊处理。

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