pyspark 中没有时间戳的滚动平均值

Question

我们可以在 pyspark 中使用 window function 找到时间序列数据的 rolling/moving 平均值。

我正在处理的数据没有任何 timestamp 列，但它确实有一个 strictly increasing 列 frame_number。数据看起来像这样。

d = [
    {'session_id': 1, 'frame_number': 1, 'rtd': 11.0, 'rtd2': 11.0,}, 
    {'session_id': 1, 'frame_number': 2, 'rtd': 12.0, 'rtd2': 6.0}, 
    {'session_id': 1, 'frame_number': 3, 'rtd': 4.0, 'rtd2': 233.0}, 
    {'session_id': 1, 'frame_number': 4, 'rtd': 110.0, 'rtd2': 111.0,}, 
    {'session_id': 1, 'frame_number': 5, 'rtd': 13.0, 'rtd2': 6.0}, 
    {'session_id': 1, 'frame_number': 6, 'rtd': 43.0, 'rtd2': 233.0}, 
    {'session_id': 1, 'frame_number': 7, 'rtd': 11.0, 'rtd2': 111.0,}]

df = spark.createDataFrame(d)

+------------+-----+-----+----------+
|frame_number|  rtd| rtd2|session_id|
+------------+-----+-----+----------+
|           1| 11.0| 11.0|         1| 
|           2| 12.0|  6.0|         1|
|           3|  4.0|233.0|         1|
|           4|110.0|111.0|         1|
|           5| 13.0|  6.0|         1|
|           6| 43.0|233.0|         1|
|           7| 11.0|111.0|         1|
+------------+-----+-----+----------+

我想在严格增加的列 frame_number 上找到列 rtd 的滚动平均值。

我正在尝试这样的事情（使用 collect_list）。

window_size=2
w = Window.partitionBy("session_id").orderBy("frame_number").rowsBetween(0, window_size)
df_lists = df.withColumn('rtd_list', F.collect_list('rtd').over(w))

+------------+-----+-----+----------+-------------------+
|frame_number|  rtd| rtd2|session_id|           rtd_list|
+------------+-----+-----+----------+-------------------+
|           1| 11.0| 11.0|         1|  [11.0, 12.0, 4.0]|
|           2| 12.0|  6.0|         1| [12.0, 4.0, 110.0]|
|           3|  4.0|233.0|         1| [4.0, 110.0, 13.0]|
|           4|110.0|111.0|         1|[110.0, 13.0, 43.0]|
|           5| 13.0|  6.0|         1| [13.0, 43.0, 11.0]|
|           6| 43.0|233.0|         1|       [43.0, 11.0]|
|           7| 11.0|111.0|         1|             [11.0]|
+------------+-----+-----+----------+-------------------+

然后应用 UDF 得到移动平均线。

windudf = F.udf( lambda v: str(np.nanmean(v)), StringType())
out = df_lists.withColumn("moving_average", windudf("rtd_list"))
+------------+-----+-----+----------+-------------------+------------------+
|frame_number|  rtd| rtd2|session_id|           rtd_list|    moving_average|
+------------+-----+-----+----------+-------------------+------------------+
|           1| 11.0| 11.0|         1|  [11.0, 12.0, 4.0]|               9.0|
|           2| 12.0|  6.0|         1| [12.0, 4.0, 110.0]|              42.0|
|           3|  4.0|233.0|         1| [4.0, 110.0, 13.0]|42.333333333333336|
|           4|110.0|111.0|         1|[110.0, 13.0, 43.0]|55.333333333333336|
|           5| 13.0|  6.0|         1| [13.0, 43.0, 11.0]|22.333333333333332|
|           6| 43.0|233.0|         1|       [43.0, 11.0]|              27.0|
|           7| 11.0|111.0|         1|             [11.0]|              11.0|
+------------+-----+-----+----------+-------------------+------------------+

上述方法的问题是它无法为 window 定义 slide duration。以上方法计算每一帧的移动平均值。在找到平均值之前，我不想将我的 window 移动一些。有什么方法可以实现吗？

Answer 1

定义window:

from pyspark.sql import functions as F

w = F.window(
    F.col("frame_number").cast("timestamp"),
    # Just example
    windowDuration="10 seconds",
    slideDuration="5 seconds",
).alias("window")

(df
    .groupBy(w, F.col("session_id"))
    .avg("rtd", "rtd2")
    .withColumn("window", F.col("window").cast("struct<start:long,end:long>"))
    .orderBy("window.start")
    .show())

# +------+----------+------------------+------------------+       
# |window|session_id|          avg(rtd)|         avg(rtd2)|
# +------+----------+------------------+------------------+
# |[-5,5]|         1|             34.25|             90.25|
# |[0,10]|         1|29.142857142857142|101.57142857142857|
# |[5,15]|         1|22.333333333333332|116.66666666666667|
# +------+----------+------------------+------------------+

另外请不要使用 collect_list 和 udf 来计算平均值。它没有任何好处，并且会严重影响性能。

pyspark 中没有时间戳的滚动平均值

Rolling average without timestamp in pyspark

apache-spark

pyspark

pyspark-sql

moving-average

sliding-window