Pandas-on-Spark的apply方法返回错误结果的问题排查

问题背景 #

需要求解如下递推方程：

x[k+1] = v[k] * y[k] + x[k]
y[k+1] = 0.01 * z[k] + y[k]
其中k的取值范围为0到N-1，初始值x[0] = 0.0，y[0] = 1.0

通过DiffEqns类实现递推逻辑，小数据集测试结果正常，但在处理比测试数据大10000倍的工业级数据集时，最终输出结果完全错误。调试发现计算阶段的打印值正确，但最终DataFrame返回无效结果或错误的历史计算值。

实现代码 #

递推计算类

class DiffEqns:
    def __init__(self):
        self.x = 0.0
        self.y = 1.0
        self.prev_second = None

    def calculate_values(self, row):
        # Calculate x and y
        self.x = row["v"] * self.y + self.x
        row["x"] = self.x
        self.y = 0.01 * row["z"] + self.y
        row["y"] = self.y
        
        print(self.x, self.y)
        print(row["x"], row["y"])
        
        return row

测试代码

import pyspark.pandas as ps
import pandas as pd

# Sample input data
data = [(0.0, 1, 10),(0.5, 1, 20),(1.0, 2, 10),(1.3, 2, 20),(1.6, 2, 30),
        (2.0, 1, 10),(2.5, 1, 20),(3.0, 2, 10),(3.3, 2, 20),(3.6, 2, 30),
        ]

# Create a pyspark.pandas DataFrame with input columns t, v, and z
df = ps.DataFrame(data, columns=["v", "z"])
        
# Initialize the calculator
de_obj = DiffEqns()

# Apply the calculate_values function to the DataFrame
df = df.apply(de_obj.calculate_values, axis=1)

print(df)

问题根源 #

1. 分布式执行的状态断裂：
   Pandas-on-Spark在处理大数据集时会自动拆分数据为多个分区，分发到不同节点/进程并行处理。你创建的de_obj实例会被序列化后发送到各个分区，但类的状态（self.x、self.y）无法在不同分区之间同步。每个分区使用独立的实例副本，导致递推计算的上下文断裂——各分区内的计算基于初始值而非全局递推的累积结果。
2. Row对象的不可变性限制：
   Spark环境中传递的row对象本质是不可变的，直接修改row["x"]和row["y"]的操作在分布式场景下无法保证正确生效。当数据跨分区处理时，修改后的row无法正确合并到最终全局结果中。

解决方案 #

方法1：使用Spark自定义聚合函数实现全局递推 #

递推计算属于依赖历史状态的累计计算，通过Spark自定义聚合函数可以在分布式环境下维护全局状态：

import pyspark.pandas as ps
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, DoubleType

class DiffEqnAggregator(F.UserDefinedAggregateFunction):
    # 输入schema：v和z列
    def inputSchema(self):
        return StructType([
            StructField("v", DoubleType()),
            StructField("z", DoubleType())
        ])
    
    # 缓冲区schema：维护当前x和y的状态
    def bufferSchema(self):
        return StructType([
            StructField("x", DoubleType()),
            StructField("y", DoubleType())
        ])
    
    # 输出schema：返回当前行的x和y结果
    def dataType(self):
        return StructType([
            StructField("x", DoubleType()),
            StructField("y", DoubleType())
        ])
    
    def deterministic(self):
        return True
    
    # 初始化缓冲区状态
    def initialize(self, buffer):
        buffer[0] = 0.0  # x初始值
        buffer[1] = 1.0  # y初始值
    
    # 逐行更新状态
    def update(self, buffer, input):
        buffer[0] = input[0] * buffer[1] + buffer[0]
        buffer[1] = 0.01 * input[1] + buffer[1]
    
    # 合并不同分区的缓冲区（需保证数据全局有序，这里简化处理）
    def merge(self, buffer1, buffer2):
        # 若数据按递推顺序分区，需基于buffer1的最终状态重新计算buffer2的所有项
        # 本示例假设数据已全局有序，若分区无序需额外处理
        pass
    
    # 返回当前缓冲区对应的结果
    def evaluate(self, buffer):
        return (buffer[0], buffer[1])

# 转换为Spark DataFrame并应用聚合函数
spark_df = ps.DataFrame(data, columns=["v", "z"]).to_spark()
spark_df = spark_df.withColumn("result", DiffEqnAggregator()(F.struct("v", "z")))\
                   .select("v", "z", F.col("result.x").alias("x"), F.col("result.y").alias("y"))

# 转换回Pandas-on-Spark DataFrame
df = ps.DataFrame(spark_df)
print(df)

方法2：使用mapGroupsWithState维护全局状态 #

针对严格按递推顺序排列的数据，可通过mapGroupsWithState维护全局递推状态：

import pyspark.pandas as ps
from pyspark.sql import functions as F
from pyspark.sql.streaming import GroupState
from pyspark.sql.types import StructType, StructField, DoubleType

# 创建带有序标识的Spark DataFrame
spark_df = ps.DataFrame(data, columns=["v", "z"]).to_spark().withColumn("seq_id", F.monotonically_increasing_id())

def update_global_state(key, rows, state):
    # 初始化状态
    if not state.exists:
        current_x = 0.0
        current_y = 1.0
    else:
        current_x, current_y = state.get
    
    # 按顺序逐行计算
    for row in sorted(rows, key=lambda r: r.seq_id):
        current_x = row.v * current_y + current_x
        current_y = 0.01 * row.z + current_y
        yield (row.v, row.z, current_x, current_y)
    
    # 更新全局状态
    state.update((current_x, current_y))

# 按虚拟分组（全局单组）处理
state_schema = StructType([
    StructField("x", DoubleType()),
    StructField("y", DoubleType())
])

result_spark_df = spark_df.groupBy(F.lit(1).alias("dummy_group"))\
                          .mapGroupsWithState(update_global_state, outputMode="append", stateSchema=state_schema)\
                          .select("_1", "_2", "_3", "_4")\
                          .toDF(["v", "z", "x", "y"])

# 转换回Pandas-on-Spark DataFrame
df = ps.DataFrame(result_spark_df)
print(df)

方法3：强制单分区处理（仅临时验证） #

若仅需临时验证逻辑正确性，可强制DataFrame为单分区，避免分布式状态问题，但此方法会丢失分布式计算优势，大数据集下性能极差：

# 强制单分区后再执行apply
df = df.repartition(1).apply(de_obj.calculate_values, axis=1)

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