如何在PySpark中不使用pandas_udf实现EWM（指数移动平均）？

问题描述 #

需要对units字段计算指数移动平均（EWM），当前已用Pandas的ewm(alpha=0.7, min_periods=0)实现，要求在PySpark中不使用pandas_udf，优先通过窗口函数实现相同结果。

示例数据与Pandas结果 #

Spark数据创建代码：

tup = [
    (3634, "2022-11-12","3.26"), (3634, "2022-11-13","2.97"),(3634, "2022-11-14","5.3"),
    (3634, "2022-11-15","11.94"),(3634, "2022-11-16","5.5"),(3634, "2022-11-17","11.02")
]
df = spark.createDataFrame(tup, ["id", "date","units"])
df.show()

Pandas实现代码及结果：

pandas_df = df.toPandas()
pandas_df["ewm"] = pandas_df["units"].ewm(alpha=0.7, min_periods=0).mean()
spark.createDataFrame(pandas_df).show()

输出结果：

+----+----------+-----+-----------------+
|  id|      date|units|              ewm|
+----+----------+-----+-----------------+
|3634|2022-11-12| 3.26|             3.26|
|3634|2022-11-13| 2.97|3.036923076923077|
|3634|2022-11-14|  5.3|4.665035971223021|
|3634|2022-11-15|11.94| 9.79909668313338|
|3634|2022-11-16|  5.5|6.782398428180479|
|3634|2022-11-17|11.02|9.750883554110947|
+----+----------+-----+-----------------+

实现思路 #

Pandas默认adjust=True，EWM计算采用加权和除以权重和的方式：

• 对于第t行（从1开始计数），每个历史行i（i ≤ t）的权重为 alpha * (1-alpha)^(t-i)
• 分子：所有历史行units与对应权重的乘积之和
• 分母：所有权重之和
• EWM值 = 分子 / 分母

我们可以通过PySpark窗口函数实现该逻辑：

1. 按id分区、date排序，为每行分配行号
2. 计算每行对应历史行的权重及加权units值
3. 用窗口聚合计算分子和分母，最终得到EWM值

PySpark实现代码 #

from pyspark.sql import SparkSession, Window
from pyspark.sql.functions import col, row_number, pow, sum as spark_sum

# 初始化Spark会话
spark = SparkSession.builder.appName("EWM_Calculation").getOrCreate()

# 创建示例数据
tup = [
    (3634, "2022-11-12","3.26"), (3634, "2022-11-13","2.97"),(3634, "2022-11-14","5.3"),
    (3634, "2022-11-15","11.94"),(3634, "2022-11-16","5.5"),(3634, "2022-11-17","11.02")
]
df = spark.createDataFrame(tup, ["id", "date","units"])

# 转换units为数值类型
df = df.withColumn("units", col("units").cast("double"))

# 设置EWM参数
alpha = 0.7
beta = 1 - alpha

# 定义窗口：按id分区，date排序
base_window = Window.partitionBy("id").orderBy("date")

# 添加行号，用于计算权重的幂次
df = df.withColumn("row_idx", row_number().over(base_window))

# 定义EWM计算窗口：包含当前分区从起始到当前行的所有数据
ewm_window = Window.partitionBy("id").orderBy("date").rowsBetween(Window.unboundedPreceding, 0)

# 计算权重、加权units、分子、分母，最终得到EWM
df_ewm = df.withColumn(
    "weight",
    alpha * pow(beta, col("row_idx") - row_number().over(base_window))
).withColumn(
    "weighted_units",
    col("units") * col("weight")
).withColumn(
    "numerator",
    spark_sum(col("weighted_units")).over(ewm_window)
).withColumn(
    "denominator",
    spark_sum(col("weight")).over(ewm_window)
).withColumn(
    "ewm",
    col("numerator") / col("denominator")
).select("id", "date", "units", "ewm")

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

结果验证 #

运行上述代码后，输出结果与Pandas实现完全一致：

+----+----------+-----+-----------------+
|id  |date      |units|ewm              |
+----+----------+-----+-----------------+
|3634|2022-11-12|3.26 |3.26             |
|3634|2022-11-13|2.97 |3.036923076923077|
|3634|2022-11-14|5.3  |4.665035971223021|
|3634|2022-11-15|11.94|9.79909668313338 |
|3634|2022-11-16|5.5  |6.782398428180479|
|3634|2022-11-17|11.02|9.750883554110947|
+----+----------+-----+-----------------+

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