PySpark无需for循环按关联表对应条件过滤行的实现方法

核心思路 #

利用Spark的分布式关联能力，通过内连接匹配阈值+条件过滤的方式实现，完全不需要遍历循环，性能远高于逐行处理。

前置准备 #

请先确保两个DataFrame的date字段均已转换为DateType类型，避免字符串格式导致比较逻辑错误，类型转换示例如下：

from pyspark.sql import functions as F
from pyspark.sql.types import DateType

df_1 = df_1.withColumn("date", F.col("date").cast(DateType()))
df_2 = df_2.withColumn("date", F.col("date").cast(DateType()))

核心实现代码 #

# 重命名df2的date字段，避免和df1的同名字段冲突
df2_with_threshold = df_2.withColumnRenamed("date", "threshold_date")

# 按condition关联两个表，过滤符合要求的记录
result_df = df_1.join(
    df2_with_threshold,
    on="condition",
    how="inner"
).filter(
    F.col("date") > F.col("threshold_date")
).select("condition", "date")

执行result_df.show()即可得到预期输出：

+---------+----------+
|condition|      date|
+---------+----------+
|        A|2018-01-02|
|        A|2018-01-03|
|        B|2018-04-06|
+---------+----------+

逻辑说明 #

1. 重命名df2的日期字段为阈值字段，避免连接后同名字段歧义
2. 按condition字段做内连接，给df1每一行匹配到对应分组的阈值日期
3. 过滤掉df1日期不大于阈值的记录，最终仅保留需要的字段即可

补充：如果df2中同一个condition存在多条记录，可先对df2按condition分组聚合得到唯一阈值后再做连接，避免连接后数据膨胀，聚合示例：

# 比如取每个condition的最大阈值作为过滤标准
df2_with_threshold = df_2.groupBy("condition").agg(F.max("date").alias("threshold_date"))

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