如何在Spark初始DataFrame内标记目标行以优化性能

问题背景 #

处理包含交易信息的大型Spark DataFrame时，需要筛选同一ID下，连续两次贷款间隔小于15天且贷款额差值≤3的所有记录。原方案通过新建new_df再执行df.subtract(new_df)筛选，该操作在大数据集上性能极差，希望改为在原DataFrame中添加标记列，直接标记符合条件的行，避免不必要的数据复制和全量比对。

原始数据集结构如下：

Name    ID     ContractDate LoanSum Status
A       ID1    2022-10-10   10      Closed 
A       ID1    2022-10-15   13      Active
A       ID1    2022-10-30   20      Active
B       ID2    2022-11-05   30      Active
C       ID3    2022-12-10   40      Closed
C       ID3    2022-12-12   43      Active
C       ID3    2022-12-19   46      Active
D       ID4    2022-12-10   10      Closed
D       ID4    2022-12-12   30      Active

现有实现代码：

from pyspark.sql import functions as f
from pyspark.sql import Window

df = spark.createDataFrame(data).toDF('Name','ID','ContractDate','LoanSum','Status')
df.show()

cols = df.columns
w = Window.partitionBy('ID').orderBy('ContractDate')

new_df = df.withColumn('PreviousContractDate', f.lag('ContractDate').over(w)) \
  .withColumn('PreviousLoanSum', f.lag('LoanSum').over(w)) \
  .withColumn('Target', f.expr('datediff(ContractDate, PreviousContractDate) < 15 and LoanSum - PreviousLoanSum <= 3')) \
  .withColumn('Target', f.col('Target') | f.lead('Target').over(w)) \
  .filter('Target == True') \
  .select(cols[0], *cols[1:])

优化方案：添加标记列替代新建DataFrame #

直接在原DataFrame上添加is_valid标记列，标记该行是否属于符合条件的记录（包括自身与前一行匹配，或作为前一行与后一行匹配），无需新建DataFrame再执行subtract操作。

修改后的代码 #

from pyspark.sql import functions as f
from pyspark.sql import Window

# 初始化原始DataFrame
df = spark.createDataFrame(data).toDF('Name','ID','ContractDate','LoanSum','Status')

# 定义窗口：按ID分区，按合同日期升序排序
w = Window.partitionBy('ID').orderBy('ContractDate')

# 在原DataFrame上添加标记列
df_with_flag = df.withColumn('prev_date', f.lag('ContractDate').over(w)) \
                 .withColumn('prev_loan', f.lag('LoanSum').over(w)) \
                 # 标记当前行与前一行是否满足条件
                 .withColumn('curr_valid', f.expr('datediff(ContractDate, prev_date) < 15 and LoanSum - prev_loan <= 3')) \
                 # 标记当前行是否是下一行的匹配前项（即下一行与当前行满足条件）
                 .withColumn('next_valid', f.lead('curr_valid').over(w)) \
                 # 最终标记：当前行自身符合条件，或作为匹配对的前一行
                 .withColumn('is_valid', f.coalesce(f.col('curr_valid'), f.lit(False)) | f.coalesce(f.col('next_valid'), f.lit(False))) \
                 # 清理临时辅助列（可选，若不需要可保留）
                 .drop('prev_date', 'prev_loan', 'curr_valid', 'next_valid')

# 后续筛选目标记录直接使用标记列
filtered_df = df_with_flag.filter(f.col('is_valid') == True)
filtered_df.show()

优化说明 #

1. 避免昂贵的subtract操作：原方案中df.subtract(new_df)需要对全量数据做比对，大数据集下会产生大量shuffle和计算开销，优化后直接在原数据上完成标记，无额外数据复制。
2. 执行计划更高效：所有计算在同一个窗口操作链中完成，Spark可以更好地优化执行逻辑，减少不必要的阶段。
3. 灵活性更高：保留全量原始数据和标记列，后续既可以筛选目标记录，也能基于全量数据做其他分析。

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