如何在PySpark DataFrame中按行条件生成address_new列？

问题场景 #

现有如下PySpark DataFrame：

name       address     result
rishi     los angeles   true
tushar    california    false
keerthi   texas         false

需要新增address_new列，规则为：

• 当result为true时，address_new复制address的值
• 当result为false时，address_new设为Null

预期结果：

name       address     result  address_new
rishi     los angeles   true   los angeles
tushar    california    false   null
keerthi   texas         false   null

推荐实现方案：使用when+otherwise #

PySpark内置的when函数可以直接实现条件判断，这种方式属于分布式矢量化操作，性能远优于逐行遍历。

代码示例 #

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

# 初始化SparkSession（如果未初始化）
spark = SparkSession.builder.appName("ConditionalColumnDemo").getOrCreate()

# 构造示例DataFrame
data = [
    ("rishi", "los angeles", True),
    ("tushar", "california", False),
    ("keerthi", "texas", False)
]
df = spark.createDataFrame(data, ["name", "address", "result"])

# 生成address_new列
df_result = df.withColumn(
    "address_new",
    when(col("result") == True, col("address")).otherwise(None)
)

# 查看结果
df_result.show()

代码说明 #

• when(col("result") == True, col("address"))：匹配result为true的行，将address的值赋值给address_new
• otherwise(None)：其他情况（即result为false）直接设为Null
• 该方案无需逐行处理，完全利用PySpark的分布式计算能力，适合大数据量场景

不推荐的逐行处理方式 #

PySpark不建议用foreach、map或自定义UDF逐行处理DataFrame，会破坏分布式计算优势，导致性能大幅下降。如果一定要模拟逐行逻辑，可参考以下UDF实现，但仅适合小数据量测试：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 定义自定义函数
def get_address_new(result, address):
    return address if result else None

# 转换为UDF
address_new_udf = udf(get_address_new, StringType())

# 应用UDF生成新列
df_result = df.withColumn("address_new", address_new_udf(col("result"), col("address")))
df_result.show()

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