基于PySpark实现Excel依赖前序行的多列计算逻辑

针对包含A-I列的Excel表格（其中F、G、H、I为固定随机值），以下是对应计算逻辑的PySpark实现，核心处理行与行之间的依赖关系（A、B列依赖前一行数据，C、D、E列依赖当前行的B列）。

实现思路 #

• 为数据添加行索引，匹配Excel的行顺序，确保能准确获取前一行数据
• 使用窗口函数lag()提取前一行的B、C、D、E值，用于计算当前行的A、B
• 按Excel公式依次计算C、D、E列，再通过迭代处理A、B列的序列依赖

完整代码 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 初始化Spark会话
spark = SparkSession.builder.appName("ExcelCalcToPySpark").getOrCreate()

# 模拟示例数据：row_num对应Excel行号，F、G为固定随机值，H、I为冗余列
data = [
    (2, 12.5, 8.3, 45, 67),  # Excel第2行
    (3, 9.1, 10.2, 32, 58),  # Excel第3行
    (4, 7.7, 11.5, 29, 61)   # Excel第4行
]
df = spark.createDataFrame(data, ["row_num", "F", "G", "H", "I"])

# 定义窗口：按行号升序，用于获取前一行数据
window_spec = Window.orderBy("row_num")

# 初始化首行（Excel第2行）的A、B值
df = df.withColumn("B", F.when(F.col("row_num") == 2, 1000).otherwise(None))
df = df.withColumn("A", F.when(F.col("row_num") == 2, None).otherwise(None))

# 计算C、D、E列（依赖当前行B值）
df = df.withColumn("C", F.col("B") * 5 + (100 / 2))
df = df.withColumn("D", F.col("B") * 5 + (10 / 2))
df = df.withColumn("E", F.col("B") * 5 + (1 / 2))

# 迭代处理后续行的A、B计算（解决递归依赖）
while df.filter(F.col("B").isNull()).count() > 0:
    # 获取前一行的历史数据
    df = df.withColumn("prev_B", F.lag("B").over(window_spec))
    df = df.withColumn("prev_C", F.lag("C").over(window_spec))
    df = df.withColumn("prev_D", F.lag("D").over(window_spec))
    df = df.withColumn("prev_E", F.lag("E").over(window_spec))
    
    # 计算当前行A列：F + G + ((前一行C+D+E)/2)
    df = df.withColumn(
        "A",
        F.when(F.col("A").isNull(), F.col("F") + F.col("G") + ((F.col("prev_C") + F.col("prev_D") + F.col("prev_E")) / 2))
        .otherwise(F.col("A"))
    )
    
    # 计算当前行B列：A + (前一行B/2)
    df = df.withColumn(
        "B",
        F.when(F.col("B").isNull(), F.col("A") + (F.col("prev_B") / 2))
        .otherwise(F.col("B"))
    )
    
    # 重新计算当前行的C、D、E列（B值已更新）
    df = df.withColumn("C", F.col("B") * 5 + (100 / 2))
    df = df.withColumn("D", F.col("B") * 5 + (10 / 2))
    df = df.withColumn("E", F.col("B") * 5 + (1 / 2))
    
    # 删除临时辅助列
    df = df.drop("prev_B", "prev_C", "prev_D", "prev_E")

# 按行号排序展示结果
df.orderBy("row_num").show(truncate=False)

# 关闭Spark会话
spark.stop()

代码关键点说明 #

• row_num字段模拟Excel的行号，保证数据处理顺序与Excel一致
• 首行A、B值按Excel规则初始化：A为NULL，B=1000
• C、D、E列直接映射Excel公式，依赖当前行的B值
• 循环迭代处理后续行的A、B：由于A依赖前一行的C/D/E，B依赖前一行的B和当前行的A，必须通过迭代填充所有行的空值，直到所有计算完成

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