如何根据另一列内容动态引用DataFrame对应列的值

问题场景 #

现有如下结构的DataFrame：

+---+---+---+---+
|ref|  A|  B|  C|
+---+---+---+---+
|  A|  2|  3|  4|
|  C|  9|  8|  7|
|  B|  5|  6|  7|
+---+---+---+---+

需要新增一列new，其值为ref列内容所指向的对应列的值，最终结果如下：

+---+
|new|
+---+
|  2|
|  7|
|  6|
+---+

由于无法提前知晓ref列的可能取值，显式的F.when链式写法不适用，需要动态实现方案。

解决方案 #

方法1：利用create_map构建映射提取

通过构建列名与列值的映射关系，直接根据ref取值提取对应值：

from pyspark.sql import functions as F

# 获取除ref外的所有目标列
target_cols = [col for col in df.columns if col != "ref"]

# 构建列名到列值的映射
col_mapping = F.create_map(*[item for pair in [(F.lit(col), F.col(col)) for col in target_cols] for item in pair])

# 新增new列
df_result = df.withColumn("new", F.map_get(col_mapping, F.col("ref")))

# 查看结果
df_result.select("new").show()

说明：该方法通过create_map将所有目标列的名称和对应值打包成一个Map类型列，再用map_get根据ref列的字符串取值直接匹配提取，逻辑直观且完全动态。

方法2：动态生成when条件结合coalesce

基于显式when写法做动态扩展，遍历目标列生成对应条件后合并：

from pyspark.sql import functions as F

target_cols = [col for col in df.columns if col != "ref"]

# 动态生成所有when条件
dynamic_conditions = [F.when(F.col("ref") == col, F.col(col)) for col in target_cols]

# 用coalesce取第一个匹配的结果
df_result = df.withColumn("new", F.coalesce(*dynamic_conditions))

# 查看结果
df_result.select("new").show()

说明：该方法自动生成所有可能的when判断逻辑，通过coalesce确保取到第一个匹配的列值，完美适配未知ref取值的场景。

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