PySpark在滑动窗口上应用UDF时出现IndexError问题求解

• UDF存在变量名笔误：定义的UDF入参名称为dt，但构造Pandas DataFrame时使用了未定义的变量loc，运行时会出现逻辑异常，部分场景下生成空结果。
• 未处理空列表边界：如果滑动窗口内没有匹配的地点记录，collect_list会返回空列表，直接取values[0]就会抛出索引越界错误。take()仅拉取前几条有效数据，不会触发边界问题，全量执行（保存、排序）时遇到空窗口就会报错。
• 滑动窗口范围定义错误：你写的rangeBetween(-days(-1),days(1))等价于rangeBetween(86400, 86400)，和预期的[-1天, +1天]范围不符，需要修正窗口边界。

方案1：修正原有UDF逻辑 #

适合需要保留Pandas UDF写法的场景：

from pyspark.sql import functions as F
from pyspark.sql.window import Window
import pandas as pd
from typing import List
from pyspark.sql.types import StringType

# 修正窗口范围：按ID分区，窗口覆盖当前日期前后各1天
days = lambda i: i * 86400
w = Window.partitionBy('ID').orderBy(F.col('date').cast('timestamp').cast('long'))\
     .rangeBetween(-days(1), days(1))

# 修正UDF逻辑，处理边界情况
@F.udf(StringType)
def get_top_loc(loc_list: List[str]) -> str:
    if not loc_list:
        return None
    # 简化频次统计逻辑
    cnt_series = pd.Series(loc_list).value_counts(ascending=False)
    return cnt_series.index[0]

# 执行计算
df_ = df_.withColumn('GOAL_Window', F.collect_list('loc').over(w))\
         .withColumn('GOAL_Loc', get_top_loc(F.col('GOAL_Window')))

方案2：使用Spark内置函数（性能更优） #

Spark 3.0+推荐用内置函数实现，避免UDF的序列化开销，运行效率更高：

df_ = df_.withColumn('GOAL_Window', F.collect_list('loc').over(w))\
         # 展开窗口内的地点列表，统计每个地点的频次
         .select('*', F.explode('GOAL_Window').alias('tmp_loc'))\
         .groupBy('ID', 'date', 'loc', 'GOAL_Window', 'tmp_loc')\
         .agg(F.count('*').alias('freq'))\
         # 按频次倒序，去重保留每个ID+日期的最高频次地点
         .orderBy(F.col('freq').desc())\
         .dropDuplicates(['ID', 'date'])\
         .select('ID', 'date', 'loc', 'GOAL_Window', F.col('tmp_loc').alias('GOAL_Loc'))\
         .orderBy('ID', 'date')

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