如何在Spark Structured Streaming事件时间窗口分组中使用UDF？

你尝试的代码运行失败，本质是Spark Structured Streaming 场景下，将内联window表达式直接传入groupBy参数时，无法直接绑定GROUPED_MAP类型的Pandas UDF，这是流式引擎的底层执行路径限制，并非代码逻辑错误。
静态批处理DataFrame场景下，内联window加groupBy后接apply的写法可正常运行；但切换为流式场景时，这类带内联窗口的分组逻辑会被引擎判定为专属的流式窗口聚合执行路径，该路径目前未实现对整组传入DataFrame的GROUPED_MAP Pandas UDF的支持，因此会直接触发校验错误拦截运行。

可以通过调整代码结构绕开引擎的校验规则，实现和内置窗口聚合完全等价的自定义UDF处理效果，具体实现逻辑如下：

• 提前通过withColumn将window表达式生成的窗口列绑定到原DataFrame上，不把window调用直接嵌在groupBy的参数列表里
• 为流式窗口添加对应的Watermark配置，满足流式引擎处理迟到数据的要求
• 将提前生成的窗口列、业务分组列作为groupBy的入参，正常调用apply方法绑定自定义UDF

可直接运行的参考代码：

from pyspark.sql.functions import window, col
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

# 自定义UDF的输出schema，仅需要定义组内处理后输出的业务字段即可
custom_output_schema = StructType([
    StructField("key", StringType()),
    StructField("avg_min", DoubleType())
])

@pandas_udf(custom_output_schema, functionType=PandasUDFType.GROUPED_MAP)
def process_group(df):
    # 保留任意自定义的组内处理逻辑
    return df

# 为流式DataFrame配置Watermark，绑定提前计算的窗口列
words_windowed = words \
    .withWatermark("timestamp", "10 minutes") \
    .withColumn(
        "window",
        window(col("timestamp"), "10 minutes", "5 minutes")
    )

# 按提前生成的窗口列、业务列分组，绑定自定义UDF执行
windowed_udf_result = words_windowed.groupBy(
    col("window"),
    col("word")
).apply(process_group)

• 上述写法的窗口计算逻辑，和官方教程中内联传入window表达式的分组效果完全一致，同样会按10分钟窗口大小、5分钟滑动步长做分组，迟到数据的处理规则也完全匹配
• 自定义UDF的输出schema不需要手动写入窗口的开始、结束时间列：引擎在分组过程中会自动将窗口键作为隐式上下文携带在结果中，最终输出的DataFrame结构和内置聚合函数的输出结构完全一致
• Spark 3.0之前的版本，流式场景下的GROUPED_MAP Pandas UDF整体不支持，需要升级版本才能正常运行上述逻辑
• 部分老版本Spark中，.apply()绑定GROUPED_MAP UDF的写法在流式场景下需要替换为调用applyInPandas方法，效果完全等价，没有功能差异

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