PySpark实现含范围与分类条件的复杂表关联需求

需求说明 #

• 输入表结构：
  • Table1：包含pol_no、version、cover字段，以及多组评级因子字段（如rf_1、rf_2...）和对应规则值字段（如rl_1、rl_2...）
  • Table2：包含version、cover（关联键），以及rating_factor（评级因子标识）、cat_value（分类值）、range_start（范围左边界）、range_end（范围右边界），还有结果字段amt、coeff、rate

匹配规则 #

1. 核心关联键：version + cover，两张表先通过这两个字段做关联
2. 评级因子匹配：Table1的rf_n需匹配Table2中rating_factor的对应标识（如rf_1对应rating_factor='rf_1'）
3. 规则值匹配逻辑：
   • 若Table2的cat_value不为空，Table1的rl_n需与cat_value完全匹配（分类匹配）
   • 若Table2的range_start/range_end不为空，Table1的rl_n需满足range_start <= rl_n <= range_end（范围匹配）
   • 支持混合场景：比如一组因子是分类匹配，另一组是范围匹配，需同时满足所有匹配条件后，将amt、coeff、rate映射回Table1对应字段

实现要求 #

• 禁止硬编码因子字段，需通过动态循环处理任意数量的评级因子组
• 适配万级数据量的性能要求
• 覆盖所有混合匹配组合（[分类,范围]、[范围,分类]等）
• 确保最终结果中amt、coeff、rate字段均为有效值

PySpark代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, expr
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

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

# 示例Table1数据（实际可替换为读取外部数据源）
data_table1 = [
    ("POL001", "V1", "COVER_A", "AGE", 35.0, "REGION", "NORTH"),
    ("POL002", "V1", "COVER_A", "AGE", 45.0, "REGION", "SOUTH"),
    ("POL003", "V2", "COVER_B", "AGE", 28.0, "REGION", "EAST")
]
schema_table1 = StructType([
    StructField("pol_no", StringType(), True),
    StructField("version", StringType(), True),
    StructField("cover", StringType(), True),
    StructField("rf_1", StringType(), True),
    StructField("rl_1", DoubleType(), True),
    StructField("rf_2", StringType(), True),
    StructField("rl_2", StringType(), True)
])
table1 = spark.createDataFrame(data_table1, schema=schema_table1)

# 示例Table2数据（实际可替换为读取外部数据源）
data_table2 = [
    ("V1", "COVER_A", "AGE", None, 30.0, 40.0, 1000.0, 0.8, 0.05),
    ("V1", "COVER_A", "AGE", None, 40.0, 50.0, 1500.0, 0.7, 0.06),
    ("V1", "COVER_A", "REGION", "NORTH", None, None, 800.0, 0.9, 0.04),
    ("V1", "COVER_A", "REGION", "SOUTH", None, None, 900.0, 0.85, 0.045),
    ("V2", "COVER_B", "AGE", None, 25.0, 30.0, 1200.0, 0.88, 0.052),
    ("V2", "COVER_B", "REGION", "EAST", None, None, 750.0, 0.92, 0.038)
]
schema_table2 = StructType([
    StructField("version", StringType(), True),
    StructField("cover", StringType(), True),
    StructField("rating_factor", StringType(), True),
    StructField("cat_value", StringType(), True),
    StructField("range_start", DoubleType(), True),
    StructField("range_end", DoubleType(), True),
    StructField("amt", DoubleType(), True),
    StructField("coeff", DoubleType(), True),
    StructField("rate", DoubleType(), True)
])
table2 = spark.createDataFrame(data_table2, schema=schema_table2)

# 步骤1：动态识别Table1中的评级因子对
rf_fields = [f for f in table1.columns if f.startswith("rf_")]
rl_fields = [f for f in table1.columns if f.startswith("rl_")]
factor_pairs = list(zip(rf_fields, rl_fields))

# 步骤2：将Table2转换为宽表，按version+cover聚合，每个评级因子的规则和结果单独成列
table2_wide = table2.groupBy("version", "cover").pivot("rating_factor").agg(
    expr("first(cat_value) as cat"),
    expr("first(range_start) as range_start"),
    expr("first(range_end) as range_end"),
    expr("first(amt) as amt"),
    expr("first(coeff) as coeff"),
    expr("first(rate) as rate")
)

# 步骤3：动态构建关联匹配条件
join_conditions = [table1.version == table2_wide.version, table1.cover == table2_wide.cover]
for rf_col, rl_col in factor_pairs:
    # 提取因子编号，对应宽表中的字段前缀
    factor_prefix = rf_col
    # 分类匹配条件：仅当分类值非空时生效
    cat_match = when(col(f"{factor_prefix}_cat").isNotNull(), col(rl_col) == col(f"{factor_prefix}_cat"))
    # 范围匹配条件：仅当范围边界非空时生效
    range_match = when(
        col(f"{factor_prefix}_range_start").isNotNull() & col(f"{factor_prefix}_range_end").isNotNull(),
        col(rl_col) >= col(f"{factor_prefix}_range_start") & col(rl_col) <= col(f"{factor_prefix}_range_end")
    )
    # 合并当前因子的匹配条件（分类/范围满足其一即可）
    join_conditions.append(cat_match | range_match)

# 步骤4：执行关联并提取结果字段
result = table1.join(table2_wide, on=join_conditions, how="inner")
# 动态生成结果字段映射
select_cols = table1.columns.copy()
for rf_col in rf_fields:
    select_cols.extend([
        col(f"{rf_col}_amt").alias(f"amt_{rf_col.split('_')[1]}"),
        col(f"{rf_col}_coeff").alias(f"coeff_{rf_col.split('_')[1]}"),
        col(f"{rf_col}_rate").alias(f"rate_{rf_col.split('_')[1]}")
    ])
result = result.select(*select_cols)

# 展示最终结果
result.show()

# 停止SparkSession
spark.stop()

代码说明 #

1. 动态识别因子：通过字段前缀匹配自动提取所有评级因子对，无需硬编码字段名
2. 宽表转换：将Table2的长表结构转为宽表，让每个评级因子的规则和结果单独成列，简化关联逻辑
3. 动态条件构建：循环每个因子对，自动生成分类/范围匹配条件，适配任意混合场景
4. 结果映射：按因子编号动态命名结果字段，确保输出结构清晰

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