Spark UDF实现Lookup：匹配col1并返回col2≤key的表A字段

我来帮你搞定这个需求！要实现基于表B的col1和col2去匹配表A中同col1且col2≤目标值的记录，用UDF的话得先做好数据预处理，不然每次查询都扫全表性能会很差。下面分两种场景给你演示（返回所有匹配记录/返回最接近的最大col2记录），顺便也给你提个不用UDF的高效方案。

第一步：准备测试数据 #

先把你给的表A和表B转换成Spark DataFrame，这里用Python代码示例（Scala逻辑完全类似）：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, broadcast, collect_list, struct
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, IntegerType

# 初始化SparkSession
spark = SparkSession.builder.appName("LookupUDFDemo").getOrCreate()

# 创建表A的DataFrame
data_a = [
    ("A", 123, "d1", "d2"),
    ("A", 134, "d3", "d4"),
    ("B", 156, "d5", "d6"),
    ("B", 178, "d7", "d8")
]
schema_a = StructType([
    StructField("col1", StringType()),
    StructField("col2", IntegerType()),
    StructField("col3", StringType()),
    StructField("col4", StringType())
])
df_a = spark.createDataFrame(data_a, schema=schema_a)

# 创建表B的DataFrame
data_b = [
    ("A", 129),
    ("A", 147),
    ("B", 199),
    ("B", 175)
]
schema_b = StructType([
    StructField("col1", StringType()),
    StructField("col2", IntegerType())
])
df_b = spark.createDataFrame(data_b, schema=schema_b)

第二步：预处理表A并广播 #

为了让UDF能快速查询，我们先把表A按col1分组，收集每个分组下的(col2, col3, col4)结构，然后广播这个分组数据（避免每个Task都重复加载表A，大幅提升性能）：

# 按col1分组，收集每组的匹配数据
a_grouped = df_a.groupBy("col1").agg(
    collect_list(struct("col2", "col3", "col4")).alias("lookup_data")
)

# 广播分组后的数据，减少重复IO
broadcast_a = broadcast(a_grouped)

场景1：返回所有符合条件的记录 #

如果需要把所有col1相同且col2≤目标值的col3、col4都返回，用这个UDF：

# 定义UDF：筛选符合条件的记录并返回(col3, col4)列表
lookup_all_udf = udf(
    lambda target_col2, lookup_list: [
        (item["col3"], item["col4"]) for item in lookup_list if item["col2"] <= target_col2
    ],
    ArrayType(StructType([
        StructField("col3", StringType()),
        StructField("col4", StringType())
    ]))
)

# 关联表B和广播数据，应用UDF
result_all = df_b.join(broadcast_a, on="col1", how="left") \
    .withColumn("matched_records", lookup_all_udf("col2", "lookup_data")) \
    .drop("lookup_data")

# 查看结果
result_all.show(truncate=False)

运行后结果会是：

+----+----+------------------------+
|col1|col2|matched_records         |
+----+----+------------------------+
|A   |129 |[{d1, d2}]              |
|A   |147 |[{d1, d2}, {d3, d4}]    |
|B   |199 |[{d5, d6}, {d7, d8}]    |
|B   |175 |[{d5, d6}]              |
+----+----+------------------------+

场景2：返回col2最大的匹配记录 #

如果只需要每个表B记录对应最大的符合条件的col2的那条记录（这种场景在业务中更常见），修改UDF如下：

# 定义UDF：返回符合条件中col2最大的(col3, col4)
lookup_max_udf = udf(
    lambda target_col2, lookup_list: 
        max(
            [item for item in lookup_list if item["col2"] <= target_col2],
            key=lambda x: x["col2"],
            default=None
        ) if lookup_list else None,
    StructType([
        StructField("col3", StringType()),
        StructField("col4", StringType())
    ])
)

# 应用UDF
result_max = df_b.join(broadcast_a, on="col1", how="left") \
    .withColumn("matched_max_record", lookup_max_udf("col2", "lookup_data")) \
    .drop("lookup_data")

result_max.show(truncate=False)

结果会是：

+----+----+----------------+
|col1|col2|matched_max_record|
+----+----+----------------+
|A   |129 |{d1, d2}        |
|A   |147 |{d3, d4}        |
|B   |199 |{d7, d8}        |
|B   |175 |{d5, d6}        |
+----+----+----------------+

额外推荐：不用UDF的高效方案 #

如果你的数据量很大，UDF可能无法被Spark优化器很好地处理，这时候可以用Join+窗口函数的方式实现，性能更优：

from pyspark.sql.window import Window
from pyspark.sql.functions import col, row_number

# 先关联表A和表B，筛选符合条件的记录
joined_df = df_b.alias("b").join(df_a.alias("a"), on="col1", how="left") \
    .filter(col("a.col2") <= col("b.col2"))

# 按表B的col1和col2分组，按表A的col2降序排序，取第一条（最大col2）
window_spec = Window.partitionBy("b.col1", "b.col2").orderBy(col("a.col2").desc())

result_window = joined_df.withColumn("rn", row_number().over(window_spec)) \
    .filter(col("rn") == 1) \
    .select("b.col1", "b.col2", "a.col3", "a.col4")

result_window.show(truncate=False)

这个方案不需要UDF，Spark能更好地做执行计划优化，适合大数据量场景。

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