基于Databricks的PySpark多变量方程组求解方案问询

云上老周

2026-7-12

基于Databricks的PySpark多变量方程组求解方案探讨

问题陈述与需求

需要开发Databricks上的PySpark程序，求解多变量递推方程组并生成x、y列：

输入：包含浮点型数据的v、z列，数据有效且条目数量一致
输出：生成x、y列，x初始值为0，y初始值为1，满足以下递推公式：
```
x[k+1] = v[k] * y[k] + x[k]
y[k+1] = 0.01 * z[k] + y[k]
```
额外规则：当t列的整数部分（秒数）变化时，x重置为0，y重置为1
最终输出需包含v、z、x、y所有列

当前方案选择与分析

因递推计算存在行依赖，无法全局并行，因此选定PySpark.Pandas作为实现方案：

对比本地Pandas，PySpark.Pandas运行在集群上，规避了本地内存与资源限制
虽无法发挥Spark分布式计算能力，但这是当前适配大数据场景的可行方案
已实现Pandas、Koalas、PySpark.Pandas三种版本，输出均符合要求

现有实现代码

Pandas 实现

import pandas as pd

# 示例输入数据
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)]

# 创建Pandas DataFrame
df = pd.DataFrame(data, columns=["t", "v", "z"])

# 初始化x、y列及状态变量
df["x"] = 0.0
df["y"] = 0.0
x = 0.0
y = 1.0
prev_second = int(df.at[0, "t"])

# 递推计算x、y
for i in df.index:
    current_second = int(df.at[i, "t"])
    
    # 秒数变化时重置状态
    if current_second > prev_second:
        x = 0.0
        y = 1.0
        prev_second = current_second
    
    # 计算当前x、y值
    x = df.at[i, "v"] * y + x
    df.at[i, "x"] = x
    y = 0.01 * df.at[i, "z"] + y 
    df.at[i, "y"] = y

# 输出结果
print(df)

Koalas 实现

import databricks.koalas as ks

# 示例输入数据
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)]

# 创建Koalas DataFrame
df = ks.DataFrame(data, columns=["t", "v", "z"])

# 封装状态与计算逻辑的类
class diffEq_ks:
    def __init__(self):
        self.x = 0.0
        self.y = 1.0
        self.prev_second = None

    def calculate_values(self, row):
        if self.prev_second is None:
            self.prev_second = int(row["t"])
        current_second = int(row["t"])
        
        # 秒数变化时重置状态
        if current_second > self.prev_second:
            self.x = 0.0
            self.y = 1.0
            self.prev_second = current_second
        
        # 计算当前x、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
        return row

# 初始化计算器并应用计算
Eq_ks = diffEq_ks()
df = df.apply(Eq_ks.calculate_values, axis=1)

# 输出结果
print(df)

PySpark.Pandas 实现（选定方案）

import pyspark.pandas as ps

# 示例输入数据
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)]

# 创建PySpark.Pandas DataFrame
df = ps.DataFrame(data, columns=["t", "v", "z"])

# 封装状态与计算逻辑的类
class diffEq_pyps:
    def __init__(self):
        self.x = 0.0
        self.y = 1.0
        self.prev_second = None

    def calculate_values(self, row):
        if self.prev_second is None:
            self.prev_second = int(row["t"])
        current_second = int(row["t"])
        
        # 秒数变化时重置状态
        if current_second > self.prev_second:
            self.x = 0.0
            self.y = 1.0
            self.prev_second = current_second
        
        # 计算当前x、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

        return row
    
# 初始化计算器并应用计算
Eq_pyps = diffEq_pyps()
df = df.apply(Eq_pyps.calculate_values, axis=1)

# 输出结果
print(df)

三种实现的统一输出

t    v     z    x    y
0  0.0  1.0  10.0  1.0  1.1
1  0.5  1.0  20.0  2.1  1.3
2  1.0  2.0  10.0  2.0  1.1
3  1.3  2.0  20.0  4.2  1.3
4  1.6  2.0  30.0  6.8  1.6
5  2.0  1.0  10.0  1.0  1.1
6  2.5  1.0  20.0  2.1  1.3
7  3.0  2.0  10.0  2.0  1.1
8  3.3  2.0  20.0  4.2  1.3
9  3.6  2.0  30.0  6.8  1.6

替代方案建议

虽然递推计算无法全局并行，但可基于t的整数部分（秒数）做分组并行处理，这是更贴合Spark原生能力的方案：

按second = int(t)分组，每组内的计算独立（组间无依赖），可分布式执行
对每个分组内的数据按t排序，确保递推顺序正确
使用Spark的pandas_udf（分组聚合UDF）实现组内递推，既能利用集群资源，又符合Spark分布式模型

Spark原生实现示例（分组并行版）

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

# 初始化SparkSession
spark = SparkSession.builder.appName("RecursiveEquation").getOrCreate()

# 定义Schema
schema = StructType([
    StructField("t", DoubleType(), True),
    StructField("v", DoubleType(), True),
    StructField("z", DoubleType(), True)
])

# 示例输入数据
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)]

# 创建Spark DataFrame
df = spark.createDataFrame(data, schema=schema)

# 添加分组键：秒数
df = df.withColumn("second", F.floor(F.col("t")).cast("integer"))

# 定义分组计算的pandas UDF，输出包含x、y的Schema
output_schema = StructType([
    StructField("t", DoubleType(), True),
    StructField("v", DoubleType(), True),
    StructField("z", DoubleType(), True),
    StructField("second", DoubleType(), True),
    StructField("x", DoubleType(), True),
    StructField("y", DoubleType(), True)
])

@pandas_udf(output_schema, functionType=F.PandasUDFType.GROUPED_MAP)
def calculate_group(group_df):
    # 分组内按t排序
    group_df = group_df.sort_values("t").reset_index(drop=True)
    # 初始化x、y列
    group_df["x"] = 0.0
    group_df["y"] = 0.0
    x = 0.0
    y = 1.0
    for i in group_df.index:
        x = group_df.loc[i, "v"] * y + x
        group_df.loc[i, "x"] = x
        y = 0.01 * group_df.loc[i, "z"] + y
        group_df.loc[i, "y"] = y
    return group_df

# 应用分组计算并整理结果
result_df = df.groupBy("second").apply(calculate_group)
result_df = result_df.drop("second").orderBy("t")

# 展示结果
result_df.show()

方案对比

PySpark.Pandas方案：代码简洁，与Pandas使用习惯一致，但apply(axis=1)本质是逐行串行处理，仅能利用单节点资源
Spark原生分组方案：通过分组实现并行计算，不同秒数的分组可在集群不同节点同时处理，充分利用Spark分布式能力，更适合超大规模数据场景

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

基于Databricks的PySpark多变量方程组求解方案问询

问题陈述与需求 #

当前方案选择与分析 #

现有实现代码 #