Apache Spark TransformWithState算子未按预期工作问题排查

我在流处理场景中实现了一个自定义StatefulProcessor，用于基于上一行数据计算丢包率，但出现异常：首次全量微批处理正常，后续微批的输出行数逐渐减少，最终仅处理前两行后就停止输出实时数据。不确定是状态无法召回还是关联逻辑有误，最终写入时会过滤空值。

实现代码 #

df_base = (
    add_oid(df_mef,"managedGroupID","terminalID")
    .withWatermark("time_trunc", "15 minutes")
)

counters = df_base.select(
    "OId", "time_trunc",
    "fwdPacketLoss", "fwdPacketsSent",
    "rtnPacketLoss", "rtnPacketsSent"
)

class PacketLossProcessor(StatefulProcessor):
    def init(self, handle: StatefulProcessorHandle) -> None:
        self.handle = handle
        state_schema = StructType([
            StructField("fwd_sent", DoubleType(), True),
            StructField("fwd_loss", DoubleType(), True),
            StructField("rtn_sent", DoubleType(), True),
            StructField("rtn_loss", DoubleType(), True)
        ])
        self.state = handle.getValueState("packet_loss_state", state_schema)

    def handleInputRows(self, key, rows: Iterator[pd.DataFrame], timerValues) -> Iterator[pd.DataFrame]:
        # ---- DEBUG START (safe) ----------------------------
        oid = key[0] if isinstance(key, (list, tuple)) else key  # safe extraction; no logic change
        # ---- DEBUG END -------------------------------------

        # Ensure state is a DataFrame
        if self.state.exists():
            state_df = self.state.get()
            if isinstance(state_df, pd.DataFrame) and not state_df.empty:
                state = state_df.iloc[0].to_dict()
            else:
                state = {
                    "fwd_sent": None, "fwd_loss": None,
                    "rtn_sent": None, "rtn_loss": None
                }
        else:
            state = {
                "fwd_sent": None, "fwd_loss": None,
                "rtn_sent": None, "rtn_loss": None
            }
        # Process input rows
        all_rows = pd.concat(list(rows)).sort_values("time_trunc")
        output = []

        for _, r in all_rows.iterrows():
            ts = r["time_trunc"]
            fs, fl = r["fwdPacketsSent"], r["fwdPacketLoss"]
            rs, rl = r["rtnPacketsSent"], r["rtnPacketLoss"]


            def pct(prev_s, prev_l, cur_s, cur_l):
                if pd.isna(prev_s) or pd.isna(prev_l) or pd.isna(cur_s) or pd.isna(cur_l): return None

                #PR For modem resets we are treating the loss as Null(currently only if the counter increases we treat it as valid)
                if cur_s <= prev_s: return None 
                ds, dl = cur_s - prev_s, cur_l - prev_l
                if ds <= 0: return None
                val = 100.0 * (dl / ds)
                return max(0.0, min(100.0, val))

            fwd_pct = pct(state["fwd_sent"], state["fwd_loss"], fs, fl)
            rtn_pct = pct(state["rtn_sent"], state["rtn_loss"], rs, rl)

            row_out = {"OId": str(oid), "time_trunc": ts}
            if fwd_pct is not None:
                row_out["fwdpacketLoss"] = fwd_pct
            if rtn_pct is not None:
                row_out["rtnpacketLoss"] = rtn_pct

            if "fwdpacketLoss" in row_out or "rtnpacketLoss" in row_out:
                output.append(row_out)


            state["fwd_sent"], state["fwd_loss"] = fs, fl
            state["rtn_sent"], state["rtn_loss"] = rs, rl

        # Update state as a DataFrame
        self.state.update(pd.DataFrame([state]))
        yield pd.DataFrame(output)

output_schema = StructType([
    StructField("OId", StringType(), True),
    StructField("time_trunc", TimestampType(), True),
    StructField("fwdpacketLoss", DoubleType(), True),
    StructField("rtnpacketLoss", DoubleType(), True)
])
    
df_loss = (
    counters.groupBy("OId")
    .transformWithStateInPandas(
        statefulProcessor=PacketLossProcessor(),
        outputStructType=output_schema,
        outputMode="append",
        timeMode="ProcessingTime"
    )
)

输入数据 #

当前输出结果 #

问题根源与修复方案 #

核心问题 #

1. 数据排序顺序错误：输入数据是时间倒序（从05:50到05:34），但代码中sort_values("time_trunc")默认升序，导致处理顺序变成05:34→05:35→...→05:50。处理完05:35后，状态中保存的是05:35的包数，后续处理05:36到05:50时，当前包数cur_s会小于状态中的prev_s（因为时间更早的数据包数更小），触发pct函数里的cur_s <= prev_s判断，直接返回None，没有输出。
2. 时间模式不匹配：设置了基于事件时间的Watermark，但transformWithStateInPandas使用timeMode="ProcessingTime"，可能导致状态清理逻辑混乱。

修复步骤 #

1. 修正数据排序方向

确保按时间从早到晚处理，将排序代码改为：

all_rows = pd.concat(list(rows)).sort_values("time_trunc", ascending=True)

或者如果输入数据是倒序，也可以直接反转：

all_rows = pd.concat(list(rows)).sort_values("time_trunc", ascending=False)

关键是保证处理顺序和时间递增一致，这样包数才会符合递增的预期。

2. 统一时间模式

如果依赖事件时间的Watermark，将timeMode改为EventTime：

df_loss = (
    counters.groupBy("OId")
    .transformWithStateInPandas(
        statefulProcessor=PacketLossProcessor(),
        outputStructType=output_schema,
        outputMode="append",
        timeMode="EventTime"
    )
)

3. 调试辅助（可选）

在处理每行时添加日志，监控状态和计算结果：

import logging
logging.basicConfig(level=logging.INFO)

# 在循环内添加
logging.info(f"OId {oid} | 时间 {ts} | 当前发送包数 {fs} | 上次发送包数 {state['fwd_sent']} | 计算结果 {fwd_pct}")

4. 优化包数校验逻辑（可选）

如果存在包数重置的场景，可增加波动阈值，避免误判：

def pct(prev_s, prev_l, cur_s, cur_l):
    if pd.isna(prev_s) or pd.isna(prev_l) or pd.isna(cur_s) or pd.isna(cur_l): 
        return None
    # 允许一定范围内的包数波动，避免重置误判
    if cur_s < prev_s - 100:  
        return None 
    ds, dl = cur_s - prev_s, cur_l - prev_l
    if ds <= 0: 
        return None
    val = 100.0 * (dl / ds)
    return max(0.0, min(100.0, val))

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