PySpark Databricks中递归计算DataFrame的A列问题求助

针对你的问题，递归计算可以通过数学转换为累乘运算，结合PySpark窗口函数实现分布式高效计算，避免逐行处理的性能瓶颈。

公式转换 #

先拆解递归逻辑：

• ( A_1 = A_0 \times (1 + \frac{B_0}{100}) )
• ( A_2 = A_1 \times (1 + \frac{B_1}{100}) = A_0 \times (1+\frac{B_0}{100}) \times (1+\frac{B_1}{100}) )
• 以此类推，第n行的( A_n = 初始A值 \times \prod_{i=0}^{n-1} (1+\frac{B_i}{100}) )

通过对数转换将累乘转为累加（PySpark无内置累乘窗口函数）：
( \log(\prod x_i) = \sum \log(x_i) )，再通过指数还原乘积：( \exp(\sum \log(x_i)) = \prod x_i )

代码实现（Databricks PySpark环境） #

from pyspark.sql.window import Window
from pyspark.sql.functions import col, row_number, exp, sum as spark_sum, log, monotonically_increasing_id

# 1. 创建示例DataFrame（替换为你的实际数据）
data = [(-15), (-5), (-10)]
df = spark.createDataFrame(data, "double").withColumnRenamed("value", "B")

# 2. 添加行号确保迭代顺序（如果已有排序列，直接用该列排序即可）
order_window = Window.orderBy(monotonically_increasing_id())
df = df.withColumn("row_num", row_number().over(order_window))

# 3. 计算每个B对应的系数
df = df.withColumn("coefficient", 1 + col("B")/100)

# 4. 计算系数的对数，转为累加操作
df = df.withColumn("log_coeff", log(col("coefficient")))

# 5. 窗口累加对数，再指数化得到累乘结果，乘以初始A值3740
cum_window = Window.orderBy("row_num").rowsBetween(Window.unboundedPreceding, Window.currentRow)
initial_A = 3740
df = df.withColumn("A", initial_A * exp(spark_sum(col("log_coeff")).over(cum_window)))

# 查看结果
df.select("row_num", "B", "A").show(truncate=False)

结果输出 #

+-------+----+------------------+
|row_num| B  |A                 |
+-------+----+------------------+
|1      |-15 |3179.0            |
|2      |-5  |3020.05           |
|3      |-10 |2718.045          |
+-------+----+------------------+

优势说明 #

• 基于窗口函数的分布式计算，适配大数据量场景，性能远优于逐行处理或递归UDF
• 无需循环迭代，完全利用Spark的并行计算能力

注意事项 #

• 确保( 1 + \frac{B}{100} )为正数，否则对数运算会报错（你的案例中B值均满足条件）
• 如果数据已有明确排序键（如时间戳、ID），替换monotonically_increasing_id()为该键即可保证顺序正确

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