Spark SQL(PySpark)如何筛选col2值存在于col1列的数据行

逐行for循环是本地单机编程的思路，不适合Spark分布式计算场景，不仅容易出逻辑错误，数据量稍大还会直接内存溢出，下面给3种可直接运行的实现方式，都能得到目标结果。

样例数据初始化 #

先构造和给出样例完全一致的测试数据集：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder.appName("col_filter").getOrCreate()

# 构造样例数据
source_data = [
    ("01011245", "ABE", "ATL"),
    ("01020600", "ABE", "DTW"),
    ("01021245", "DTW", "ATL"),
    ("01020605", "ABE", "ATL")
]
df = spark.createDataFrame(source_data, schema=["date", "col1", "col2"])

方法1：左半连接实现（生产环境/大数据量首选） #

左半连接是Spark里专门用来做存在性判断的连接方式，不会产生重复数据，全程在分布式计算层完成，不需要把全量col1值拉到Driver节点内存，性能最优。

# 提取col1列所有去重值作为匹配参照
col1_unique = df.select(col("col1").alias("match_key")).distinct()

# 左半连接过滤符合条件的行
result = df.join(
    col1_unique,
    on=df.col2 == col1_unique.match_key,
    how="left_semi"
)

result.show()

运行输出：

+--------+----+----+
|    date|col1|col2|
+--------+----+----+
|01020600| ABE| DTW|
+--------+----+----+

方法2：isin函数实现（小数据集场景写法简洁） #

如果col1列去重后值的数量不多（万级以内），可以把去重值收集为本地列表，用isin函数直接过滤，写法更简单。如果col1去重后值超过十万级别不要用这个方法，会把大量数据拉到Driver端造成内存压力，序列化传输列表也会增加开销。

# 收集col1去重值为Python列表
col1_val_list = df.select("col1").distinct().rdd.flatMap(lambda row: row).collect()

# 过滤
result = df.filter(col("col2").isin(col1_val_list))
result.show()

方法3：纯Spark SQL实现 #

如果习惯写SQL逻辑，直接用IN子查询即可，Spark会自动做底层优化，性能和左半连接一致：

# 注册临时视图
df.createOrReplaceTempView("source_table")

# 执行SQL过滤
result = spark.sql("""
SELECT * FROM source_table
WHERE col2 IN (SELECT DISTINCT col1 FROM source_table)
""")
result.show()

for循环失效原因 #

Spark DataFrame是分区块存储在集群多个节点上的分布式数据集，不支持直接像本地Pandas对象那样逐行遍历做判断。硬写for循环要么是没有触发Action计算拿不到实际值，要么是把全量数据拉到Driver节点本地循环，完全丢失分布式计算的能力，数据量稍微大一点就会直接OOM报错。

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