PySpark如何高效统计两个DataFrame中符合区间要求的日期数量

性能瓶颈原因 #

你当前的实现性能差的核心原因是每循环一行deadline就触发一次count() action，有多少条deadline数据就会执行多少次Spark作业，大量的作业调度和重复扫描数据的开销会随着数据量上涨变得不可接受。

优化方案 #

方案1：原生Spark API实现（通用性最强，无额外依赖） #

全程使用Spark原生算子，走Catalyst优化器优化，仅触发一次作业，适合绝大多数场景：

from pyspark.sql import functions as F

def count_days_opt(deadlines_df, dates_df, ref_date):
    # 提前过滤小于参考日期的无效数据，减少后续计算量
    valid_dates_df = dates_df.filter(F.col("DT_DATE") >= ref_date)
    # 新增常量连接键用于关联
    deadlines_df = deadlines_df.withColumn("join_key", F.lit(1))
    valid_dates_df = valid_dates_df.withColumn("join_key", F.lit(1))
    # 左关联后分组统计符合条件的日期数
    res_df = deadlines_df.join(valid_dates_df, on="join_key", how="left") \
        .groupBy("DEADLINES") \
        .agg(
            F.sum(F.when(F.col("DT_DATE") <= F.col("DEADLINES"), 1).otherwise(0)).cast("int").alias("dt_count")
        )
    return res_df

# 调用示例
new_df = count_days_opt(deadlines_df, dates_df, "2021-03-31")

如果过滤后的valid_dates_df数据量较小（小于10万条），可以将关联逻辑修改为deadlines_df.join(F.broadcast(valid_dates_df), on="join_key", how="left")，通过广播小表进一步提升性能。

方案2：广播+二分查找优化（适合日期量远小于deadline量的场景） #

如果DT_DATE的量级远低于DEADLINE的量级，可以用广播有序数组+二分查找的方案，仅触发两次作业，性能更高：

from pyspark.sql import functions as F
import bisect

def count_days_broadcast_opt(deadlines_df, dates_df, ref_date):
    # 收集有效日期并排序
    valid_dates = sorted(
        [row.DT_DATE for row in dates_df.filter(F.col("DT_DATE") >= ref_date).collect()]
    )
    # 广播有序数组到所有executor
    bc_valid_dates = spark.sparkContext.broadcast(valid_dates)
    # 定义UDF用二分查找快速统计符合条件的数量
    @F.udf(returnType="int")
    def count_qualified(deadline):
        return bisect.bisect_right(bc_valid_dates.value, deadline)
    
    return deadlines_df.withColumn("dt_count", count_qualified(F.col("DEADLINES")))

# 调用示例
new_df = count_days_broadcast_opt(deadlines_df, dates_df, "2021-03-31")

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