PySpark DataFrame高效过滤：保留两列交集行的方法

问题描述 #

我有一个PySpark DataFrame，需要过滤出满足以下条件的行：

• 该行col_a的值存在于整个DataFrame的col_b列的所有值中
• 该行col_b的值存在于整个DataFrame的col_a列的所有值中

示例DataFrame定义如下：

df = spark.createDataFrame(
    [
        ("abc", "ddc", 1, 4.5),
        ("abb", "ddc", 4, 9.1),
        ("baa", "abc", 2, 3.2),
        ("abb", "bca", 1, 5.1),
        ("ddc", "abc", 2, 3.6),
        ("abc", "baa", 3, 2.6)
    ],
    ["col_a", "col_b", "col_c", "col_d"]
)

原实现方式通过将数据拉取到Driver端生成交集集合，再用isin过滤，但这种方法在大数据集下效率极低，还会触发Spark警告：

col_a_and_col_b = (
    set(df.select("col_a").distinct().toPandas()["col_a"].tolist())
    .intersection(df.select("col_b").distinct().toPandas()["col_b"].tolist())
)

filtered_df = df.filter(
    (col("col_a").isin(col_a_and_col_b))
    & (col("col_b").isin(col_a_and_col_b))
)

期望得到的过滤结果：

+-----+-----+-----+-----+
|col_a|col_b|col_c|col_d|
+-----+-----+-----+-----+
|  abc|  ddc|    1|  4.5|
|  baa|  abc|    2|  3.2|
|  ddc|  abc|    2|  3.6|
|  abc|  baa|    3|  2.6|
+-----+-----+-----+-----+

解决方案（Spark原生高效实现） #

方法一：使用exists子查询 #

利用Spark的exists函数构造子查询条件，无需将数据拉取到Driver端，Spark会自动优化执行计划：

from pyspark.sql.functions import col, exists

# 构造条件1：col_a的值存在于col_b的唯一值集合中
cond_a_in_b = exists(df.select("col_b").distinct(), lambda b: b.col_b == col("col_a"))
# 构造条件2：col_b的值存在于col_a的唯一值集合中
cond_b_in_a = exists(df.select("col_a").distinct(), lambda a: a.col_a == col("col_b"))

# 应用过滤条件
filtered_df = df.filter(cond_a_in_b & cond_b_in_a)

方法二：使用内连接过滤 #

通过两次内连接分别过滤符合条件的行，基于Spark分布式执行，避免Driver端数据瓶颈：

# 提取col_b的唯一值并改名，避免列名冲突
b_unique = df.select("col_b").distinct().withColumnRenamed("col_b", "valid_val")
# 提取col_a的唯一值并改名
a_unique = df.select("col_a").distinct().withColumnRenamed("col_a", "valid_val")

# 第一步：过滤col_a在col_b集合中的行
temp_df = df.join(b_unique, df.col_a == b_unique.valid_val, how="inner").drop("valid_val")
# 第二步：过滤col_b在col_a集合中的行
filtered_df = temp_df.join(a_unique, temp_df.col_b == a_unique.valid_val, how="inner").drop("valid_val")

# 保留原DataFrame的列顺序（可选）
filtered_df = filtered_df.select(df.columns)

两种方法都能高效处理大数据集，避免了将大量数据拉取到Driver端的性能问题。

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