如何在PySpark中高效查询匹配值对的行？

我有一个包含id1、id2、value列的超大型数据表，id1和id2各自不唯一，但两者的组合是唯一的。现有一个包含唯一(id1,id2)值对的PySpark DataFrame，需要获取这些值对对应的value：

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("session") \
    .getOrCreate()

data = [
    ("a1", "Z1"),
    ("b1", "Y1"),
    ("a1", "Y1"),
    ("c1", "X1")
]

pyspark_df = spark.createDataFrame(data, ["id1", "id2"])

我最初的实现方式如下：

id1_list = [row.id1 for row in pyspark_df.select("id1").collect()]
id2_list = [row.id2 for row in pyspark_df.select("id2").collect()]

def get_values(id1_list, id2_list):
    import pyspark.sql.functions as F
    table = spark.table("table")
    table = table.filter(F.col("id1").isin(id1_list))
    table = table.filter(F.col("id2").isin(id2_list))
    return table

get_values(id1_list, id2_list)

但这种方式会获取所有id1在列表中且id2在列表中的行，而非指定的(id1,id2)唯一值对，不符合需求。需要在超大型数据库中实现高效的精准查询。

直接通过Inner Join实现(id1,id2)值对的精准匹配，这是最符合需求且高效的方式：

基础实现 #

import pyspark.sql.functions as F

# 加载超大型数据表
large_table = spark.table("table")

# 执行Inner Join，仅保留匹配的(id1,id2)组合
result = large_table.join(
    pyspark_df,
    on=["id1", "id2"],
    how="inner"
).select("id1", "id2", "value")

性能优化：广播Join（适用于pyspark_df数据量较小的场景） #

如果存储目标值对的pyspark_df数据量不大，可以使用广播Join，将小表广播到所有计算节点，避免超大型表的shuffle操作，大幅提升查询效率：

from pyspark.sql.functions import broadcast

large_table = spark.table("table")

result = large_table.join(
    broadcast(pyspark_df),
    on=["id1", "id2"],
    how="inner"
).select("id1", "id2", "value")

原方法不可行的原因 #

原代码中两个isin的逻辑是：只要id1在列表中且id2在列表中就会被筛选，不关心这两个值是否来自同一个指定的配对。比如若原表存在(a1, X1)，而id1_list包含a1、id2_list包含X1，这条无关数据也会被错误选中，无法实现精准匹配指定值对的需求。

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