PySpark结构化流窗口聚合缺失值前向填充技术问询

要解决指定窗口内tag无数据时，用前一窗口该tag的平均值填补的问题，我们可以通过生成全量窗口-tag组合 + 窗口函数拉取前值的方式实现，以下是修改后的完整代码：

# Python built-in
import tempfile
from datetime import datetime

# Third-party
from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import window, col, avg, coalesce, lag
from pyspark.sql.types import (
    StructType,
    StructField,
    StringType,
    TimestampType,
)
from pyspark.sql.window import Window

SCHEMA = StructType([
    StructField("timestamp", TimestampType(), True),
    StructField("tag", StringType(), True),
    StructField("value", StringType(), True),
    ]
)
ROWS = [
    Row(timestamp=datetime(year=2023, month=12, day=30, hour=0, minute=37),
        tag="A",
        value="2.8"),
    Row(timestamp=datetime(year=2023, month=12, day=30, hour=0, minute=37),
        tag="B",
        value="9.0"),
    Row(timestamp=datetime(year=2023, month=12, day=30, hour=0, minute=38),
        tag="B",
        value="7.4"),
    Row(timestamp=datetime(year=2023, month=12, day=30, hour=0, minute=47),
        tag="A",
        value="3.2"),
    Row(timestamp=datetime(year=2023, month=12, day=30, hour=0, minute=48),
        tag="A",
        value="1.8"),
]

spark = (SparkSession.builder
         .appName("structured streaming windowed aggregation")
         .getOrCreate()
)

# Suppress most log messages
spark.sparkContext.setLogLevel('ERROR')

with tempfile.TemporaryDirectory() as d:

    # Write input data to JSON file to be stream read
    df = spark.createDataFrame(ROWS, SCHEMA)
    df.printSchema()
    df.show()
    df.write.mode('overwrite').format('json').save(d)

    # Stream read the data
    df_streaming = (
        spark.readStream
        .format('json')
        .schema(SCHEMA)
        .load(d)
    )

    # 1. 原窗口聚合逻辑（新增数值类型转换）
    df_windowed_average = (
        df_streaming
        .withWatermark("timestamp", '0 minutes')
        .groupBy(window("timestamp", '15 minutes').alias("window"), col("tag"))
        .agg(avg(col("value").cast("double")).alias("avg"))
    )

    # 2. 生成所有窗口与所有tag的全量组合
    all_windows = df_windowed_average.select("window").distinct()
    all_tags = df_windowed_average.select("tag").distinct()
    full_window_tag = all_windows.crossJoin(all_tags)

    # 3. 左连接原聚合结果，保留所有窗口-tag组合
    df_full = full_window_tag.join(
        df_windowed_average,
        on=["window", "tag"],
        how="left"
    )

    # 4. 用前一窗口的tag平均值填补当前窗口的null值
    w = Window.partitionBy("tag").orderBy(col("window.start"))
    df_filled = df_full.withColumn(
        "avg",
        coalesce(col("avg"), lag(col("avg"), 1).over(w))
    ).orderBy("window", "tag")

    # Print result to terminal
    query = (
        df_filled.writeStream
        .format('console')
        .outputMode('complete')
        .trigger(once=True)
        .start()
        .awaitTermination()
    )

关键逻辑说明 #

1. 数值类型转换：原value字段是StringType，计算平均值前需转为double类型，避免类型错误。
2. 全量组合生成：通过crossJoin生成所有窗口与tag的笛卡尔积，确保每个窗口下每个tag都有记录，不会因无数据丢失。
3. 前值填充：用lag窗口函数按tag分组、窗口时间排序，拉取前一个窗口的平均值；再用coalesce优先使用当前窗口的计算值，无值时则用前一窗口的值填补。

执行结果 #

运行修改后的脚本后，输出会包含第二个窗口的tag B的平均值：

+-------------------+---+-----+
|          timestamp|tag|value|
+-------------------+---+-----+
|2023-12-30 00:17:00|  A|  2.8|
|2023-12-30 00:17:00|  B|  9.0|
|2023-12-30 00:18:00|  B|  7.4|
|2023-12-30 00:31:00|  A|  3.2|
|2023-12-30 00:32:00|  A|  1.8|
+-------------------+---+-----+

-------------------------------------------
Batch: 0
-------------------------------------------
+--------------------+---+---+
|              window|tag|avg|
+--------------------+---+---+
|{2023-12-30 00:15...|  A|2.8|
|{2023-12-30 00:15...|  B|8.2|
|{2023-12-30 00:30...|  A|2.5|
|{2023-12-30 00:30...|  B|8.2|
+--------------------+---+---+

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