基于另一列的列名数组生成对应值列的Spark实现问题

示例输入 #

from pyspark.sql import Row
import pyspark.sql.functions as F

input_df = spark.createDataFrame([
    Row(id=123, alert=1, operation=1, lookup=[]),
    Row(id=234, alert=0, operation=0, lookup=['alert']),
    Row(id=345, alert=1, operation=0, lookup=['operation']),
    Row(id=456, alert=0, operation=1, lookup=['alert', 'operation']),
])

预期输出 #

已尝试方案及问题 #

报错的尝试代码 #

input_df.withColumn("lookup_values", F.transform(F.col("lookup"), lambda x: input_df[f'{x}'])).show()

执行后报错：

AnalysisException: [UNRESOLVED_COLUMN.WITH_SUGGESTION] A column or function parameter with the name Column<'x_1'> cannot be resolved. Did you mean one of the following? [id, alert, operation, lookup].

无报错但不符合预期的代码 #

input_df.withColumn("lookup_values", F.transform(F.col("lookup"), lambda x: input_df['alert'])).show()

输出结果：

解决方案 #

问题核心是transform中的lambda参数x是列内的字符串值，无法直接用input_df[x]引用DataFrame的列，需要用Spark内置函数动态解析列名对应的数值。

方法1：transform结合expr动态解析列名 #

input_df.withColumn(
    "lookup_values",
    F.transform(
        F.col("lookup"),
        lambda col_name: F.expr(f"`{col_name}`")
    )
).show(truncate=False)

F.expr可以将字符串形式的列名解析为对应列的数值，反引号用于兼容列名含特殊字符的场景。

方法2：构建字段映射表后取值 #

适合需要查找的字段较多的场景：

# 定义所有可能需要查询的字段
target_columns = ["alert", "operation"]

# 先构建字段名到字段值的映射Map
input_df.withColumn(
    "column_map",
    F.map_from_arrays(
        F.array(*[F.lit(col) for col in target_columns]),
        F.array(*[F.col(col) for col in target_columns])
    )
# 从映射表中提取lookup列指定的字段值
).withColumn(
    "lookup_values",
    F.transform(F.col("lookup"), lambda x: F.col("column_map")[x])
).drop("column_map").show(truncate=False)

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