Spark Scala如何使用UDF更新DataFrame的Struct类型列子字段值

Spark DataFrame的Struct类型列本身是不可变的，无法直接修改子字段，需要通过UDF返回全新的Struct对象来完成更新。

首先定义UDF的返回类型，和原有colStruct的结构完全一致：

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

# 返回结构和原colStruct的字段顺序、类型、可空属性保持一致
return_schema = StructType([
    StructField("subCol1", IntegerType(), nullable=True),
    StructField("subCol2", StringType(), nullable=True),
    StructField("subCol3", IntegerType(), nullable=True)
])

编写UDF自定义更新逻辑，示例中我们将subCol1翻倍、subCol3加10，subCol2保持原值不变：

@F.udf(returnType=return_schema)
def update_struct(col_struct):
    # 空值判断避免运行时报错
    if col_struct is None:
        return None
    new_sub1 = col_struct.subCol1 * 2
    new_sub3 = col_struct.subCol3 + 10
    # 返回的元组顺序要和return_schema的字段顺序完全匹配
    return (new_sub1, col_struct.subCol2, new_sub3)

调用UDF替换原有的colStruct列即可：

df = df.withColumn("colStruct", update_struct(F.col("colStruct")))

逻辑和PySpark一致，语法略有不同：
首先定义返回的Struct Schema：

import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row

val returnSchema = new StructType()
  .add("subCol1", IntegerType, nullable = true)
  .add("subCol2", StringType, nullable = true)
  .add("subCol3", IntegerType, nullable = true)

定义UDF：

val updateStructUdf = udf((colStruct: Row) => {
  if (colStruct == null) null
  else {
    val sub1 = colStruct.getAs[Int]("subCol1")
    val sub2 = colStruct.getAs[String]("subCol2")
    val sub3 = colStruct.getAs[Int]("subCol3")
    // 自定义更新逻辑
    Row(sub1 * 2, sub2, sub3 + 10)
  }
}, returnSchema)

调用替换列：

val updatedDf = df.withColumn("colStruct", updateStructUdf(col("colStruct")))

• UDF返回的Struct的字段顺序、类型、可空属性必须和原Struct完全一致，否则会报类型不匹配错误
• 必须做好原Struct列的空值判断，避免运行时空指针异常
• 如果你使用的是Spark 3.1+版本，也可以不用UDF，直接用内置的withField方法更新Struct子字段，语法更简洁：

# Spark 3.1+ 内置方法实现（仅做扩展参考）
df = df.withColumn("colStruct", 
    F.col("colStruct")
    .withField("subCol1", F.col("colStruct.subCol1") * 2)
    .withField("subCol3", F.col("colStruct.subCol3") + 10)
)

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