Databricks中SQL循环引用问题：求Spark或SQL替代实现方案

计算需求 #

需要基于基础数据计算以下3列，存在循环依赖：

• Points = BasePoints - BurnRateCumul
• BurnRate = Points / DaysLeft
• BurnRateCumul = 截至当前行的BurnRate累计值

基础数据 #

已尝试方案及问题 #

• 使用Databricks窗口函数，无法处理循环依赖，结果不符合预期，代码如下：

Select
  bb.Date,
  bb.BasePoints,
  bb.DaysLeft,
  bb.BasePoints - sum(0) Over( order by bb.Date) AS Points,
  (bb.BasePoints - sum(0) Over( order by bb.Date)) / bb.DaysLeft AS BurnRate,
  sum(0) Over( order by bb.Date) AS BurnRateCumul
from
  burnbase as bb
order by bb.Date

• 尝试T-SQL递归CTE，但Databricks SQL不支持该语法
• 尝试压缩数据换角度处理，仍无法解决循环引用问题

期望结果 #

额外要求 #

最后一行需满足(BurnRateCumul + BurnRate) = BasePoints，且最后一天保留可燃烧数值（不为0）。

解决方案 #

方案1：Spark逐行处理（Python） #

由于计算依赖前一行的BurnRateCumul，可以用Spark的mapPartitions维护累计状态，逐行计算：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, DateType, IntegerType, DoubleType

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

# 定义结果Schema
result_schema = StructType([
    StructField("Date", DateType(), True),
    StructField("BasePoints", IntegerType(), True),
    StructField("Points", DoubleType(), True),
    StructField("DaysLeft", IntegerType(), True),
    StructField("BurnRate", DoubleType(), True),
    StructField("BurnRateCumul", DoubleType(), True)
])

# 读取基础数据并按Date排序
df = spark.table("burnbase").orderBy("Date")

# 自定义分区处理函数，维护累计BurnRateCumul
def calculate_burn_rate(partition):
    cumul_burn = 0.0
    for row in partition:
        date = row.Date
        base_points = row.BasePoints
        days_left = row.DaysLeft
        
        points = base_points - cumul_burn
        burn_rate = points / days_left
        cumul_burn += burn_rate
        
        # 保留三位小数对齐期望结果
        yield (date, base_points, round(points, 3), days_left, round(burn_rate, 3), round(cumul_burn, 3))

# 应用函数并转换为DataFrame
result_df = df.rdd.mapPartitions(calculate_burn_rate).toDF(result_schema)

# 展示结果
result_df.show()

方案2：Spark SQL迭代计算（规避循环引用） #

如果偏好SQL，可以通过迭代方式逐步计算累计值，适合数据量不大的场景：

-- 初始化临时表，添加初始BurnRateCumul为0
CREATE OR REPLACE TEMP VIEW burn_initial AS
SELECT 
    Date,
    BasePoints,
    DaysLeft,
    0.0 AS BurnRateCumul
FROM burnbase
ORDER BY Date;

-- 迭代计算第1天
CREATE OR REPLACE TEMP VIEW burn_iter1 AS
SELECT
    Date,
    BasePoints,
    DaysLeft,
    ROUND((BasePoints - BurnRateCumul)/DaysLeft + BurnRateCumul, 3) AS BurnRateCumul
FROM burn_initial
WHERE Date = '2023-01-01';

-- 迭代计算第2天
CREATE OR REPLACE TEMP VIEW burn_iter2 AS
SELECT * FROM burn_iter1
UNION ALL
SELECT
    b.Date,
    b.BasePoints,
    b.DaysLeft,
    ROUND((b.BasePoints - (SELECT BurnRateCumul FROM burn_iter1 WHERE Date = DATE_SUB(b.Date, 1)))/b.DaysLeft + (SELECT BurnRateCumul FROM burn_iter1 WHERE Date = DATE_SUB(b.Date, 1)), 3) AS BurnRateCumul
FROM burn_initial b
WHERE Date = '2023-01-02';

-- 重复上述逻辑直到覆盖所有日期，最后计算Points和BurnRate
SELECT
    Date,
    BasePoints,
    ROUND(BasePoints - BurnRateCumul, 3) AS Points,
    DaysLeft,
    ROUND((BasePoints - BurnRateCumul)/DaysLeft, 3) AS BurnRate,
    BurnRateCumul
FROM burn_iter8; -- 替换为最后一次迭代的视图

注意：迭代SQL方案需要根据数据行数手动调整迭代次数，适合小数据集；Spark RDD逐行处理方案更通用，适合大数据量场景。

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