如何在Azure Databricks的PySpark循环中创建动态自定义函数实现行复制与金额递增

核心实现思路 #

基于Spark分布式算子完成全流程逻辑，避免逐行循环带来的性能与稳定性问题，核心规则拆解：

• 单条原始记录的总生成行数 = Interval * Cycle
• 逐行的日期偏移：基于原始日期，按行顺序累加1个月
• 金额递增规则：每满一个完整Cycle周期（即Interval行），BaseAmount按给定增长率上调，每周期的涨幅按累计复利计算
• 带%标识的增长率字段需提前转换为可计算的数值型小数

具体代码实现 #

步骤1：字段预处理 #

完成增长率格式转换、总行数预计算：

from pyspark.sql import functions as F
from pyspark.sql.types import IntegerType, DoubleType, DateType
# 清洗增长率字段：移除%符号，转换为小数数值，如"2%" -> 0.02
df = df.withColumn(
    "rate_val",
    F.regexp_replace(F.col("RateIncrease(%)"), "%", "").cast(DoubleType()) / 100
)
# 预计算单条记录需生成的总行数
df = df.withColumn(
    "total_rows",
    F.col("Interval") * F.col("Cycle")
)

步骤2：序列展开实现行复制 #

为每条记录生成从0到总行数-1的连续序号，通过explode展开为多行：

df_expanded = df.withColumn(
    "row_seq",
    F.expr("sequence(0, total_rows - 1)")
).select(
    "*",
    F.explode("row_seq").alias("row_offset")
)

步骤3：逐行计算业务值 #

按行偏移量计算对应日期、所属周期、调整后的BaseAmount：

df_final = df_expanded.withColumn(
    # 计算当前行所属的周期序号：每Interval行归为一个周期，从0开始计数
    "cycle_group",
    (F.col("row_offset") / F.col("Interval")).cast(IntegerType())
).withColumn(
    # 计算滚动后的日期：基于原始日期加对应行数的月数偏移
    "calc_date",
    F.add_months(F.col("Date").cast(DateType()), F.col("row_offset"))
).withColumn(
    # 计算对应周期的BaseAmount：按复利累计涨幅，四舍五入取整
    "calc_amount",
    F.round(
        F.col("BaseAmount") * F.pow(1 + F.col("rate_val"), F.col("cycle_group"))
    ).cast(IntegerType())
).select(
    # 对齐目标表结构输出字段
    F.col("OPID"),
    F.col("calc_date").alias("Date"),
    F.col("calc_amount").alias("BaseAmount"),
    F.col("Interval"),
    F.col("Cycle"),
    F.col("RateIncrease(%)")
)

结果校验匹配 #

• 针对OPID为O1的原始记录（Interval=12，Cycle=2）：
  • 共生成24行
  • 前12行（row_offset 0-11）：周期序号为0，BaseAmount保持4375，日期从2014-07-27逐月累加至2015-06-27
  • 后12行（row_offset 12-23）：周期序号为1，BaseAmount为round(4375*1.02)=4463，日期从2015-07-27逐月累加至2016-06-27，与示例输出完全一致
• 针对OPID为O2的原始记录（Interval=7，Cycle=3）：
  • 共生成21行
  • 每7行触发一次5%的金额递增，日期从2020-12-23逐月累加

性能说明 #

该实现全程基于Spark内置SQL算子执行，无Python侧逐行循环逻辑，可支撑千万级原始记录的批量处理，不会出现大循环导致的Driver端OOM问题。

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