PySpark技术求助：如何为数组内每个Struct元素新增statusrating字段

需求说明 #

需要为DataFrame中form_info数组的每个Struct元素新增statusrating字段，映射规则如下：

• 当status为"APPROVED"时，statusrating取值1
• 当status为"PENDING"时，statusrating取值0
• 当status为"REJECTED"时，statusrating取值-1

当前DataFrame结构 #

form_info:array
|   element:struct
|   |    id :string
|   |    status:string

期望DataFrame结构 #

form_info:array
|   element:struct
|   |    id :string
|   |    status:string
|   |    statusrating:integer

解决方案代码 #

from pyspark.sql import functions as F

# 对form_info数组进行转换，为每个struct元素新增statusrating字段
df = df.withColumn(
    "form_info",
    F.transform(
        "form_info",
        lambda elem: F.struct(
            elem["id"],
            elem["status"],
            F.when(elem["status"] == "APPROVED", 1)
             .when(elem["status"] == "PENDING", 0)
             .when(elem["status"] == "REJECTED", -1)
             .alias("statusrating")
        )
    )
)

# 可选：打印Schema验证结果
df.printSchema()

代码说明 #

• 使用F.transform遍历form_info数组的每个Struct元素
• 通过F.struct将原有的id、status字段与新生成的statusrating字段重新组合为新的Struct
• 利用F.when实现类似SQL CASE WHEN的逻辑，完成status到整数的映射

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