如何用Databricks SQL或Python实现累计差值计算？

问题背景 #

有如下数据表（按rn排序），需更新abc列：

• 首行值：bpc_balance_distance_km - journey_distance
• 后续每行值：上一行的abc结果 - 当前行的journey_distance

原MySQL查询使用变量实现，需转换为Databricks SQL支持的写法，同时提供Python实现方案。

原始数据表

方案1：Databricks SQL实现 #

Databricks SQL支持窗口函数实现累积计算，无需使用变量，直接通过累积求和完成需求：

基础版（单rake_name场景） #

WITH cte AS (
    SELECT 
        *,
        SUM(journey_distance) OVER (ORDER BY rn) AS total_journey
    FROM derived_abc
)
SELECT 
    *,
    ROUND(bpc_balance_distance_km - total_journey) AS final_result
FROM cte;

分组版（多rake_name场景） #

如果存在多个rake_name，需按分组计算：

WITH cte AS (
    SELECT 
        *,
        SUM(journey_distance) OVER (PARTITION BY rake_name ORDER BY rn) AS total_journey
    FROM derived_abc
)
SELECT 
    *,
    ROUND(bpc_balance_distance_km - total_journey) AS final_result
FROM cte;

逻辑说明

• 首行：bpc_balance_distance_km - journey_distance（累积总和等于第一行的journey_distance）
• 后续行：bpc_balance_distance_km - 前n行journey_distance总和，等价于上一行结果减去当前行journey_distance，完全匹配需求。

方案2：Python（PySpark）实现 #

方法1：累积求和（推荐） #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import sum, round, col

spark = SparkSession.builder.appName("abc_calculation").getOrCreate()

# 定义窗口规则：按rake_name分组，按rn排序
window_spec = Window.partitionBy("rake_name").orderBy("rn")

# 计算累积行程距离，再推导最终结果
df = df.withColumn("total_journey", sum(col("journey_distance")).over(window_spec))
df = df.withColumn("final_result", round(col("bpc_balance_distance_km") - col("total_journey")))

# 查看结果
df.show()

方法2：lag()函数迭代计算 #

模拟逐行迭代逻辑，通过lag()获取上一行结果：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import lag, when, round, col

spark = SparkSession.builder.appName("abc_calculation").getOrCreate()

window_spec = Window.partitionBy("rake_name").orderBy("rn")

# 先计算首行值，再用lag获取上一行结果迭代计算
df = df.withColumn(
    "final_result",
    when(
        col("rn") == 1,
        col("bpc_balance_distance_km") - col("journey_distance")
    ).otherwise(
        lag(col("final_result")).over(window_spec) - col("journey_distance")
    )
)

# 对结果取整
df = df.withColumn("final_result", round(col("final_result")))

# 查看结果
df.show()

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