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

时间:2018-04-30 17:44:31

标签: python apache-spark pyspark apache-spark-sql

我正在尝试使用Grouped 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的滞后。

1 个答案:

答案 0 :(得分: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|
# +----+----+-----+------------------+

我不会使用太多的熊猫,所以可能会有更优雅的方式来做这件事。