我们可以用 when òr UDF 在spark中根据情况插入新的列。
创建示例Dataframe：

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('AddConditionalColumn').getOrCreate()

data = [(1,"baz",0),(2,"baz",0),(3,"333",2),(4,"444",1)]

columns = ["id","foo","bar"]
df = spark.createDataFrame(data = data, schema = columns)
df.show()

+---+---+---+
| id|foo|bar|
+---+---+---+
|  1|baz|  0|
|  2|baz|  0|
|  3|333|  2|
|  4|444|  1|
+---+---+---+

使用时间：

from pyspark.sql.functions import when
df2 = df.withColumn("foo", when(df.bar == 2,"X")
                           .when(df.bar == 1,"Y")
                           .otherwise(df.foo))
df2.show()

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

使用自定义项：

import pyspark.sql.functions as F
from pyspark.sql.types import *

def executeRule(value):
    if value == 2:
        return 'X'
    elif value == 1:
        return 'Y'
    else:
        return value

# Converting function to UDF

ruleUDF = F.udf(executeRule, StringType())
df3 = df.withColumn("foo", ruleUDF("bar"))
df3.show()

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

when 可以与其他列一起使用。你可以用 F.col 要获取另一列的值并提供适当的条件，请执行以下操作：

import pyspark.sql.functions as F

df2 = df.withColumn(
    'foo', 
    F.when(F.col('bar') == 2, 'X')
     .when(F.col('bar') == 1, 'Y')
     .otherwise(F.col('foo'))
)

df2.show()
+---+---+---+
| id|foo|bar|
+---+---+---+
|  1|baz|  0|
|  2|baz|  0|
|  3|  X|  2|
|  4|  Y|  1|
+---+---+---+