如何在PySpark中根据另一列的两个条件设置目标列的值

pyspark.sql.functions.when 完全支持基于其他列的取值做条件判断，实现目标列的自定义值设置，是处理这类场景的标准PySpark方法。

实现代码 #

首先导入所需函数：

from pyspark.sql.functions import when, col

使用withColumn重写foo列值，嵌套条件判断逻辑：

updated_df = original_df.withColumn(
    "foo",
    when(col("bar") == 2, "X")
    .when(col("bar") == 1, "Y")
    .otherwise(col("foo"))
)

逻辑匹配说明 #

• 第一个判断分支：匹配bar列值等于2的行，将对应foo列设置为X
• 第二个判断分支：匹配bar列值等于1的行，将对应foo列设置为Y
• 兜底分支：未被前述条件命中的行，直接保留foo列的原有取值

处理后输出示例 #

对应原始输入数据，处理完成后的DataFrame结果为：

+---+----+---+
|id |foo |bar|
+---+----+---+
|1  |baz |0  |
|2  |baz |0  |
|3  |X   |2  |
|4  |Y   |1  |
+---+----+---+

注：when方法的两个入参均支持动态表达式，判断条件可以引用其他列的值，返回值也可以直接传入列对象、运算表达式，并非仅能处理静态固定值。

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