如何在PySpark的CASE WHEN中动态应用YAML存储的列级条件

YAML条件定义 #

Name:
  criteria_1:
    column_1 = 'A' and
    column_2 = 'AA' and
    column_3 = 'AAA' and
    column_2 != column_3 and
    column_1 - column_2 <= n

  criteria_2:
    column_1 = 'B' and
    column_2 = 'BB' and
    column_3 = 'BBB'

当前尝试的代码 #

def get_criteria(criteria):
    conds = yaml.load(conds)
    try:
        cons = yaml["Name"][criteria]
        return cons
    except Exception as e:
        pass
    return None

spark.udf.register("extract_conditions", get_criteria, StringType())

df = spark.table("table")
df.createOrReplaceTempView("df")
spark.sql("select extract_conditions(df.criteria) as cds, (case when \
           (extract_conditions(df.criteria)) then 'Ok' \
           else 'Reject' end) \
           as outcome from df").show(1,False)

错误信息 #

[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "CASE WHEN extract_conditions(criteria) THEN Ok ELSE Reject END" due to data type mismatch: Parameter 1 requires the "BOOLEAN" type, however "extract_conditions(criteria)" has the type "STRING".; line 1 pos 7;

硬编码可行示例 #

select extract_conditions(df.criteria) as cds, (case when \
           (column_1 = 'A' and column_2 = 'AA' and column_3 = 'AAA') then 'Ok' \
           else 'Reject' end) \
           as outcome from df

执行结果：

+-------------------------------------------------------+-------+
|cds                                                    |outcome|
+-------------------------------------------------------+-------+
|column_1 = 'A' and column_2 = 'AA' and column_3 = 'AAA'|Ok     |
|column_1 = 'A' and column_2 = 'AA' and column_3 = 'AAA'|Ok     |
+-------------------------------------------------------+-------+

示例数据 #

已尝试的失败方案 #

• 使用eval和字符串转换
• 用字典存储条件字符串并根据criteria列取值调用

解决方案 #

方法1：预生成CASE WHEN语句（推荐） #

Spark SQL是静态解析的，无法在运行时动态解析字符串形式的条件，因此可以提前加载YAML中的所有条件，拼接成完整的CASE WHEN SQL语句后执行。

import yaml

# 加载并解析YAML条件
conds_str = '''
Name:
  criteria_1:
    column_1 = 'A' and
    column_2 = 'AA' and
    column_3 = 'AAA' and
    column_2 != column_3 and
    column_1 - column_2 <= n

  criteria_2:
    column_1 = 'B' and
    column_2 = 'BB' and
    column_3 = 'BBB'
'''
conds = yaml.safe_load(conds_str)
criteria_map = conds['Name']

# 清理条件字符串（去除换行、多余空格），并拼接CASE WHEN分支
case_clauses = []
cleaned_conds = {}
for crit_name, condition in criteria_map.items():
    clean_cond = ' '.join([line.strip() for line in condition.split('\n') if line.strip()])
    cleaned_conds[crit_name] = clean_cond
    case_clauses.append(f"WHEN criteria = '{crit_name}' AND ({clean_cond}) THEN 'Ok'")

# 组装完整的CASE WHEN表达式
case_stmt = f"CASE {' '.join(case_clauses)} ELSE 'Reject' END"

# 生成获取条件字符串的子查询
cond_str_subquery = f"""
(SELECT cond_str FROM (VALUES {', '.join([f"('{k}', '{v}')" for k, v in cleaned_conds.items()])} ) AS t(crit, cond_str) WHERE t.crit = df.criteria)
"""

# 执行最终SQL
df = spark.table("table")
df.createOrReplaceTempView("df")
spark.sql(f"""
SELECT 
  {cond_str_subquery} AS cds,
  {case_stmt} AS outcome
FROM df
""").show(truncate=False)

方法2：使用DataFrame API结合expr动态解析条件 #

如果偏好DataFrame API而非SQL，可以用expr函数将条件字符串解析为Spark表达式，再通过when匹配对应的criteria：

import yaml
from pyspark.sql.functions import expr, col, when

# 加载YAML条件
conds_str = '''
Name:
  criteria_1:
    column_1 = 'A' and
    column_2 = 'AA' and
    column_3 = 'AAA' and
    column_2 != column_3 and
    column_1 - column_2 <= n

  criteria_2:
    column_1 = 'B' and
    column_2 = 'BB' and
    column_3 = 'BBB'
'''
conds = yaml.safe_load(conds_str)
criteria_map = conds['Name']

# 清理条件字符串
cleaned_criteria = {k: ' '.join([line.strip() for line in v.split('\n') if line.strip()]) for k, v in criteria_map.items()}

# 加载原始数据
df = spark.table("table")

# 为每个criteria生成对应的布尔条件列
for crit_name, cond in cleaned_criteria.items():
    df = df.withColumn(f"cond_{crit_name}", expr(cond))

# 构建outcome列：匹配criteria并对应到布尔条件
outcome_expr = when(col("criteria") == "criteria_1", col("cond_criteria_1"))
for crit_name in list(cleaned_criteria.keys())[1:]:
    outcome_expr = outcome_expr.when(col("criteria") == crit_name, col(f"cond_{crit_name}"))
outcome_expr = outcome_expr.otherwise(False)

# 添加cds列（显示对应条件字符串）和outcome列
df = df.withColumn("cds", when(col("criteria") == "criteria_1", cleaned_criteria["criteria_1"])
                   .when(col("criteria") == "criteria_2", cleaned_criteria["criteria_2"])
                   .otherwise(None))
df = df.withColumn("outcome", when(outcome_expr, "Ok").otherwise("Reject"))

# 查看结果
df.select("cds", "outcome").show(truncate=False)

核心原因说明 #

Spark SQL在解析阶段会严格校验数据类型，UDF返回的字符串无法被当作布尔条件使用——Spark不会在运行时动态解析字符串为SQL表达式（这既存在安全风险，也不符合Spark的静态编译优化逻辑）。因此必须提前将条件整合到SQL语句或DataFrame表达式中，而非通过UDF返回字符串动态拼接。

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