PySpark DataFrame按条件动态替换列值的实现方法

问题描述 #

原始DataFrame结构与数据：

需求：当is_column字段值为1时，将Cols字段中的列名替换为对应列的数值；当is_column为0时，保留Cols原有值，最终得到如下结果：

解决方案 #

方法一：使用when结合expr动态引用列 #

这种方法直接通过表达式动态解析Cols中的列名，获取对应值，代码简洁直观：

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

# 初始化SparkSession
spark = SparkSession.builder.appName("replace_col_val").getOrCreate()

# 构建原始DataFrame
data = [
    (4, 6, 2, "A1", 1),
    (5, 1, 3, "A3", 1),
    (1, 3, 4, "22", 0)
]
df = spark.createDataFrame(data, ["A1", "A2", "A3", "Cols", "is_column"])

# 按条件替换Cols字段
df_result = df.withColumn(
    "Cols",
    when(df.is_column == 1, expr("`${Cols}`")).otherwise(df.Cols)
)

# 查看结果
df_result.show()

说明：expr("${Cols}")会将每行Cols的值作为列名，动态取出该列的数值，反引号用于兼容列名包含特殊字符的场景。

方法二：使用映射表批量处理（适合多列场景） #

如果需要替换的列较多，可以先构建列名与对应值的映射表，再通过映射表取值：

from pyspark.sql import SparkSession
from pyspark.sql.functions import map_from_arrays, lit, array, col, when

spark = SparkSession.builder.appName("replace_col_val").getOrCreate()

data = [
    (4, 6, 2, "A1", 1),
    (5, 1, 3, "A3", 1),
    (1, 3, 4, "22", 0)
]
df = spark.createDataFrame(data, ["A1", "A2", "A3", "Cols", "is_column"])

# 构建列名数组和对应值数组
target_cols = ["A1", "A2", "A3"]
col_names = array(*[lit(col_name) for col_name in target_cols])
col_values = array(*[col(col_name) for col_name in target_cols])

# 创建映射表并替换Cols字段
df_result = df.withColumn(
    "col_mapping",
    map_from_arrays(col_names, col_values)
).withColumn(
    "Cols",
    when(df.is_column == 1, col("col_mapping").getItem(col("Cols"))).otherwise(col("Cols"))
).drop("col_mapping")  # 移除临时映射表列

df_result.show()

说明：这种方法通过map_from_arrays将目标列的名称和值打包成键值对映射，后续直接通过Cols的值从映射中提取对应数值，无需为每个列单独写判断逻辑。

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