python - 如何在pyspark中运行指数加权移动平均线

Question

我正在尝试使用分组 map Pandas UDF 在 PySpark 中运行指数加权移动平均值。但它不起作用:

def ExpMA(myData):

    from pyspark.sql.functions import pandas_udf
    from pyspark.sql.functions import PandasUDFType
    from pyspark.sql import SQLContext 

    df = myData
    group_col = 'Name'
    sort_col = 'Date'

    schema = df.select(group_col, sort_col,'count').schema
    print(schema)

    @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
    def ema(pdf):
        Model = pd.DataFrame(pdf.apply(lambda x: x['count'].ewm(span=5, min_periods=1).mean()))
        return Model

    data = df.groupby('Name').apply(ema)

    return data

我也尝试在没有 Pandas udf 的情况下运行它，只是在 PySpark 中编写 ewma 方程，但问题是 ewma 方程包含当前 ewma 的滞后。

Answer 1

首先，您的 Pandas 代码不正确。不管 Spark 与否，这都是行不通的

pdf.apply(lambda x: x['count'].ewm(span=5, min_periods=1).mean())

另一个问题是输出架构，根据您的数据，它不会真正适应结果:

• 如果想添加 ewm 架构应该进行扩展。
• 如果您只想返回 ewm，则架构太大。
• 如果您只想替换，它可能与类型不匹配。

假设这是第一个场景(我允许自己稍微重写一下代码):

from pyspark.sql.functions import pandas_udf
from pyspark.sql.functions import PandasUDFType
from pyspark.sql.types import DoubleType, StructField

def exp_ma(df, group_col='Name', sort_col='Date'):
    schema = (df.select(group_col, sort_col, 'count')
        .schema.add(StructField('ewma', DoubleType())))

    @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
    def ema(pdf):
        pdf['ewm'] = pdf['count'].ewm(span=5, min_periods=1).mean()
        return pdf

    return df.groupby('Name').apply(ema)

df = spark.createDataFrame(
    [("a", 1, 1), ("a", 2, 3), ("a", 3, 3), ("b", 1, 10), ("b", 8, 3), ("b", 9, 0)], 
    ("name", "date", "count")
)

exp_ma(df).show()
# +----+----+-----+------------------+                                            
# |Name|Date|count|              ewma|
# +----+----+-----+------------------+
# |   b|   1|   10|              10.0|
# |   b|   8|    3| 5.800000000000001|
# |   b|   9|    0|3.0526315789473686|
# |   a|   1|    1|               1.0|
# |   a|   2|    3|               2.2|
# |   a|   3|    3| 2.578947368421052|
# +----+----+-----+------------------+

我不使用太多 Pandas，所以可能有更优雅的方法来做到这一点。