使用 PySpark 动态生成 CASE WHEN 语句

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

# 创建 SparkSession
spark = SparkSession.builder.appName("dynamic_case_when").getOrCreate()

# 示例数据
map_data = [('a', 'b', 'c', 'good'), ('a', 'a', '*', 'very good'),
          ('b', 'd', 'c', 'bad'), ('a', 'b', 'a', 'very good'),
          ('c', 'c', '*', 'very bad'), ('a', 'b', 'b', 'bad')]

columns = ["col1", "col2", 'col3', 'result']

mapping_table = spark.createDataFrame(map_data, columns)


data = [('a', 'b', 'c'), ('a', 'a', 'b'),
        ('c', 'c', 'a'), ('c', 'c', 'b'),
        ('a', 'b', 'b'), ('a', 'a', 'd')]

columns = ["col1", "col2", 'col3']
df = spark.createDataFrame([data], columns=columns)
df = df.selectExpr(*(f"col('{col}')[0] as {col}" for col in columns))

# 动态生成 CASE WHEN 语句
ressql = 'case '
for m in map_data:
    p = [f"{c} = '{v}'" for c, v in zip(columns, m[:3]) if v != "*"]
    ressql = ressql + ' when ' + ' and '.join(p) + f" then '{m[3]}'"
ressql = ressql + ' end'

# 应用 CASE WHEN 语句
df_with_result = df.withColumn('result', expr(ressql))

# 显示结果
df_with_result.show()

# 停止 SparkSession
spark.stop()

+----+----+----+---------+
|col1|col2|col3|   result|
+----+----+----+---------+
|   a|   b|   c|     good|
|   a|   a|   b|very good|
|   c|   c|   a| very bad|
|   c|   c|   b| very bad|
|   a|   b|   b|      bad|
|   a|   a|   d|very good|
+----+----+----+---------+