Pyspark实现交易行按日展开并计算客户每日净余额

核心逻辑为预处理交易数据后生成客户维度的连续日期序列，关联交易后通过窗口函数累加计算余额，比全量复制交易行的方案性能更高。

步骤1：交易数据预处理 #

• 统一客户ID大小写，将字符串格式的交易日期转换为日期类型
• 按交易类型生成带符号的交易值：Deposit计为正，withdraw计为负

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

# 初始化SparkSession
spark = SparkSession.builder.appName("customer_daily_balance").getOrCreate()

# 样例原始交易数据，实际使用可替换为读取你的数据源逻辑
raw_data = [
    ("c1", "T1", "01-10-2021", "Deposit", 500),
    ("C1", "T2", "01-15-2021", "withdraw", 100),
    ("C1", "T3", "01-30-2021", "withdraw", 300),
    ("C1", "T4", "02-20-2021", "Deposit", 200),
    ("c2", "T1", "01-01-2021", "Deposit", 500),
    ("C2", "T2", "01-15-2021", "withdraw", 100),
    ("C3", "T1", "01-30-2021", "Deposit", 300),
    ("C3", "T2", "02-20-2021", "withdraw", 200)
]
df = spark.createDataFrame(raw_data, schema=["customerID", "TransactionID", "TransactionDate", "TransactionType", "TransactionAmount"])

# 预处理交易数据
trans_df = df \
    .withColumn("customerID", F.upper(F.col("customerID"))) \
    .withColumn("TransactionDate", F.to_date(F.col("TransactionDate"), "MM-dd-yyyy")) \
    .withColumn("trans_val", F.when(F.col("TransactionType") == "Deposit", F.col("TransactionAmount"))
                .otherwise(-F.col("TransactionAmount")))

步骤2：生成每个客户的连续自然日序列 #

• 统计每个客户的首笔交易日期，结合当前日期生成每个客户需要计算的日期范围
• 生成连续日期基准表，和客户ID关联得到全量「客户+日期」的基础计算表

# 获取每个客户的首交易日、当前日期
customer_date_range = trans_df.groupBy("customerID") \
    .agg(F.min("TransactionDate").alias("first_trans_date")) \
    .withColumn("current_date", F.current_date())

# 生成足够长的连续日期基准序列，max_days可根据业务时间跨度调整
max_days = 10000
date_seq = spark.range(max_days) \
    .withColumn("CalendarDay", F.date_add(F.to_date(F.lit("2020-01-01")), F.col("id").cast("int")))

# 关联得到每个客户从首交易日到当前日的所有自然日
customer_calendar = customer_date_range.join(date_seq, 
    (F.col("CalendarDay") >= F.col("first_trans_date")) & 
    (F.col("CalendarDay") <= F.col("current_date")),
    "inner"
).select("customerID", "CalendarDay")

步骤3：关联交易计算每日净余额 #

• 连续日期表左关联单日交易汇总数据，无交易的日期交易值记为0
• 用窗口函数按客户维度按日期累加，得到每日净余额

# 先聚合计算每个客户每日的总交易值
daily_trans_sum = trans_df.groupBy("customerID", "TransactionDate") \
    .agg(F.sum("trans_val").alias("daily_trans"))

# 关联连续日期表
daily_trans = customer_calendar.join(
    daily_trans_sum,
    (customer_calendar.customerID == daily_trans_sum.customerID) & 
    (customer_calendar.CalendarDay == daily_trans_sum.TransactionDate),
    "left"
).select(customer_calendar["customerID"], "CalendarDay", "daily_trans")

# 计算累计净余额，无交易的日期自动继承前一日余额
window_spec = Window.partitionBy("customerID").orderBy("CalendarDay").rowsBetween(Window.unboundedPreceding, 0)
result_df = daily_trans.withColumn("Netbalance", F.sum(F.coalesce("daily_trans", F.lit(0))).over(window_spec))

# 可选：将日期转换为样例中的字符串格式
result_df = result_df.withColumn("CalendarDay", F.date_format("CalendarDay", "MM-dd-yyyy"))

# 输出结果
result_df.orderBy("customerID", "CalendarDay").show()

补充说明 #

如果需要输出样例中带计算过程的余额说明（如400(500-100)），可以额外加字段拼接计算逻辑，根据业务需求自行扩展即可。

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