如何在PySpark中实现依赖上一行输出的逐行状态变更运算

这类计算逻辑需要同时依赖相邻行的输入值、以及上一行的输出结果，属于典型的带状态序列计算，PySpark下有两类成熟的落地方案：

方案1：窗口预计算+分区内状态遍历（适用批处理场景） #

是生产环境最常用的轻量实现方式，步骤如下：

• 首先给数据集添加全局稳定的排序字段，比如自增行号、事件时间戳，必须保证排序唯一无冲突，避免计算顺序错乱
• 用lead()窗口函数预拉取每行对应的下一行输入值，窗口直接按排序字段orderBy即可，数据量极大的场景可搭配范围分区，保证相邻数据落到同一个分区
• 预计算每行和对应下一行输入值的众数，存为中间列
• 最后通过mapPartitions遍历预计算完成的数据集，维护临时变量存储上一行的输出结果，逐行计算当前行的最终输出

示例代码参考：

from pyspark.sql import SparkSession
from pyspark.sql.functions import lead
from pyspark.sql.window import Window
import pandas as pd
from scipy.stats import mode

spark = SparkSession.builder.appName("state_calc").getOrCreate()
# 读取原始数据集，已包含全局唯一排序字段row_id
df = spark.read.parquet("your_data_path")

# 预拉取下一行输入值
win = Window.orderBy("row_id")
df = df.withColumn("next_input", lead("input", 1).over(win))

# 预计算每行+下一行输入的众数
def calc_mode(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["current_mode"] = pdf.apply(
        lambda x: mode([x["input"], x["next_input"]], keepdims=True)[0][0], 
        axis=1
    )
    return pdf

df_pre = df.mapInPandas(
    calc_mode, 
    schema="row_id long, input long, next_input long, current_mode long"
).orderBy("row_id")

# 遍历维护前序输出状态，计算最终结果
def state_calc(iterator):
    last_output = None
    for row in iterator:
        # 首行自定义初始化逻辑
        if last_output is None:
            current_output = row["current_mode"]
        else:
            # 替换为你自己的计算逻辑：众数 + 上一行输出
            current_output = row["current_mode"] + last_output
        last_output = current_output
        yield (row["row_id"], row["input"], current_output)

result_rdd = df_pre.rdd.sortBy(lambda x: x["row_id"]).mapPartitions(state_calc)
result_df = result_rdd.toDF(["row_id", "input", "final_output"])

方案2：Structured Streaming状态API（适用实时流/增量场景） #

如果是流数据或者需要增量计算的场景，可以直接用Spark结构化流的状态算子实现：

• 用flatMapGroupsWithState算子，无分组场景可以设置全局固定的分组键，状态中存储上一行的输出结果
• 输入流按排序字段排序后，每处理一条数据先计算和下一条缓存数据的众数，再结合状态中存储的上一行输出计算当前结果，完成后更新状态即可

注意事项 #

• 无论用哪种方案，都必须保证数据的全局排序规则稳定唯一，避免顺序变动导致结果错误
• 超大数据量场景不要直接做全局排序，可以按排序字段做范围分区，保证相邻数据在同一个分区内，再分区内独立维护状态，避免全量数据shuffle到单节点

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