Pyspark中如何按指定列值拼接字符串生成新列？

原始数据 #

需求说明 #

根据Col1的取值，拼接生成新列NewCol1：

• 当Col1为Ora时，格式为300|Col2数值
• 当Col1为APP时，格式为100|Col2数值

期望结果：

• 300|1234
• 300|0123
• 100|0005
• 300|7778
• 100|9999

你的代码错误分析 #

你写的代码存在几个语法和逻辑问题：

1. 方法名拼写错误：withColunm 应该是 withColumn
2. 条件判断错误：F.col('Col1' == Ora) 语法完全错误，正确写法是 F.col('Col1') == 'Ora'（注意字符串值要加引号）
3. 拼接方式错误：不能直接把多个字符串参数传给when/otherwise，Spark需要用concat函数来拼接字符串

正确实现代码 #

方法一：使用when条件分支拼接 #

from pyspark.sql import functions as F

# 生成新列
df = df.withColumn(
    'NewCol1',
    F.when(F.col('Col1') == 'Ora', F.concat(F.lit('300|'), F.col('Col2')))
     .when(F.col('Col1') == 'APP', F.concat(F.lit('100|'), F.col('Col2')))
)

方法二：使用字典映射（更易扩展） #

如果后续有更多Col1的取值需要处理，用字典映射更方便：

from pyspark.sql import functions as F

# 定义前缀映射字典
prefix_map = {'Ora': '300|', 'APP': '100|'}
# 把字典转成Spark的map结构
spark_map = F.create_map(*[F.lit(item) for sublist in prefix_map.items() for item in sublist])

# 生成新列
df = df.withColumn(
    'NewCol1',
    F.concat(spark_map[F.col('Col1')], F.col('Col2'))
)

验证结果 #

运行上述代码后，NewCol1列会生成你期望的300|xxx或100|xxx格式的值。

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