Spark 3.2.1实现按客户分组计算调整天数的高效方法

核心需求回顾 #

每个Cust_ID分组内：

• 首行Adjusted Days固定为0
• 后续行的Adjusted Days = max(0, 上一行Fill_days + 上一行Adjusted Days)

实现方案（Python示例） #

前提：确定分组内的行顺序 #

首先必须给每个Cust_ID的行指定明确的排序规则（比如业务日期、创建时间等），否则递推结果无意义。这里假设用row_number()生成分组内的行号来确定顺序：

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

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

# 假设你的原始DataFrame名为df，包含Cust_ID和Fill_days列
window_spec = Window.partitionBy("Cust_ID").orderBy("业务排序列")  # 替换为实际排序字段，如"order_date"
df = df.withColumn("row_num", F.row_number().over(window_spec))

方案1：递归CTE（Spark 3.0+支持） #

通过递归公共表表达式逐行递推计算，适合分组内行数不多的场景：

# 定义基础CTE：取每个分组的第一行，Adjusted_Days设为0
base_df = df.filter(F.col("row_num") == 1).withColumn("Adjusted_Days", F.lit(0))
base_df.createOrReplaceTempView("base_table")

# 定义待递归的后续行
recursive_df = df.filter(F.col("row_num") > 1)
recursive_df.createOrReplaceTempView("recursive_table")

# 执行递归计算
result_df = spark.sql("""
    WITH RECURSIVE cte AS (
        SELECT Cust_ID, Fill_days, row_num, 0 AS Adjusted_Days
        FROM base_table
        UNION ALL
        SELECT curr.Cust_ID, curr.Fill_days, curr.row_num,
               CASE WHEN prev.Fill_days + prev.Adjusted_Days < 0 THEN 0 ELSE prev.Fill_days + prev.Adjusted_Days END AS Adjusted_Days
        FROM recursive_table curr
        JOIN cte prev ON curr.Cust_ID = prev.Cust_ID AND curr.row_num = prev.row_num + 1
    )
    SELECT Cust_ID, Fill_days, Adjusted_Days
    FROM cte
    ORDER BY Cust_ID, row_num
""").drop("row_num")

方案2：flatMapGroupsWithState（更高效的分布式处理） #

通过分组遍历的方式处理每个客户的行，避免递归JOIN的性能损耗，适合大数据量场景：

from pyspark.sql import Row

def process_cust_group(cust_id, rows_iterator):
    # 按行号排序，确保递推顺序正确
    sorted_rows = sorted(rows_iterator, key=lambda x: x.row_num)
    adjusted_days = 0
    result_rows = []
    for row in sorted_rows:
        # 生成当前行的结果
        result_rows.append(Row(
            Cust_ID=cust_id,
            Fill_days=row.Fill_days,
            Adjusted_Days=adjusted_days
        ))
        # 计算下一行的Adjusted_Days
        next_val = row.Fill_days + adjusted_days
        adjusted_days = 0 if next_val < 0 else next_val
    return iter(result_rows)

# 按Cust_ID分组后处理
result_df = df.rdd.groupBy(lambda x: x.Cust_ID) \
                 .flatMap(lambda x: process_cust_group(x[0], x[1])) \
                 .toDF()

注意事项 #

• 务必指定业务排序字段，不能依赖Spark默认的行顺序，否则递推结果会出错。
• 若分组内单客户行数极多（如百万级），优先选择flatMapGroupsWithState方案，递归CTE可能因多次JOIN导致性能下降。
• Scala版本的实现逻辑一致，仅语法略有不同，可将上述Python代码转换为Scala实现。

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