PySpark如何根据同一行其他列的值动态选取对应列的内容

实现思路 #

核心是将每行categoryName的字符串值作为键，匹配同名列的数值作为新增score列的结果，下面提供两种适用不同场景的实现方式：

方法1：固定分类场景（分类数量少） #

使用when条件分支逐一匹配，写法简单直观：

from pyspark.sql import functions as F

# 直接按条件匹配对应列
df_result = df.withColumn("score",
    F.when(F.col("categoryName") == "catA", F.col("catA"))
    .when(F.col("categoryName") == "catB", F.col("catB"))
    .otherwise(None)
)

方法2：动态分类场景（分类数量多/不固定） #

使用create_map构造列名到列值的映射，自动适配所有cat开头的列，无需手动枚举：

from pyspark.sql import functions as F

# 提取所有cat开头的分类列
cat_columns = [col for col in df.columns if col.startswith("cat")]

# 构造 列名->列值 的映射表达式
map_expression = F.create_map(
    *[item for col_name in cat_columns for item in [F.lit(col_name), F.col(col_name)]]
)

# 按categoryName从映射中取对应值
df_result = df.withColumn("score", map_expression[F.col("categoryName")])

两种方法执行后得到的df_result就和你给出的预期结果完全一致。

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