Spark 3 Scala UDF中数组数据的GenericRowWithSchema类型转换异常及字段修改问题

问题背景 #

需要编写Spark代码对数组字段中的code属性进行掩码处理（设置为"XXXX"），但在使用UDF时遇到类型转换和不可变对象修改的问题。

数据格式（Parquet存储，JSON示例） #

{"conditions":{"list":[{"element":{"code":"1234","category":"ABC"}},{"element":{"code":"4550","category":"EDC"}}]}}

定义的样例类 #

case class MyClass(conditions: Seq[MyItem])
case class MyItem(code: String, category: String)

现有Spark代码 #

val data = Seq(MyClass(conditions = Seq(MyItem("1234", "ABC"), MyItem("4550", "EDC"))))
import spark.implicits._

val rdd = spark.sparkContext.parallelize(data)
val ds = rdd.toDF().as[MyClass]

val maskedConditions: Column = updateArray.apply(col("conditions"))

ds.withColumn("conditions", maskedConditions)
  .select("conditions")
  .show(2)

遇到的问题 #

尝试的UDF代码无法正确修改数组元素：

def updateArray = udf((arr: Seq[MyItem]) => {
    for (i <- 0 to arr.size - 1) {
      val a = arr(i).asInstanceOf[org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema]
      val a = arr(i)
      println(a.getAs[MyItem](0))

      // 无法修改code字段
      // a.code = "XXXX"
    }
    arr
  })

• 保留类型转换行时，Spark的GenericRowWithSchema是不可变对象，无法直接修改字段值；
• 移除转换行后，抛出类型转换错误：

Caused by: java.lang.ClassCastException: org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema cannot be cast to MyItem

解决方案 #

方法1：使用Spark内置函数（推荐） #

利用Spark内置的transform和struct函数，无需编写UDF，性能更优：

import org.apache.spark.sql.functions.{transform, struct, lit, col}

val resultDs = ds.withColumn(
  "conditions",
  transform(
    col("conditions"),
    item => struct(lit("XXXX").alias("code"), item.getField("category").alias("category"))
  )
)

resultDs.select("conditions").show(false)

说明：transform遍历数组每个元素，通过struct重新构造结构体，将code固定为"XXXX"，保留原category值。

方法2：使用Dataset的map算子（类型安全） #

基于样例类的类型特性，直接用map算子处理，代码更直观且类型安全：

val maskedDs = ds.map { myClass =>
  val maskedItems = myClass.conditions.map(item => item.copy(code = "XXXX"))
  myClass.copy(conditions = maskedItems)
}

maskedDs.select("conditions").show(false)

说明：Scala样例类是不可变对象，通过copy方法生成新的MyItem实例，替换原code值，再构造新的MyClass对象。

方法3：正确编写UDF（仅作参考） #

如果必须使用UDF，需将输入类型改为Seq[Row]（Spark会将结构体序列化为Row），并返回符合结构的新Row：

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

val updateArray = udf((arr: Seq[Row]) => {
  arr.map { row =>
    // 按MyItem的字段顺序构造新Row：code为"XXXX"，保留原category
    Row("XXXX", row.getAs[String]("category"))
  }
})

val resultDs = ds.withColumn("conditions", updateArray(col("conditions")))
resultDs.select("conditions").show(false)

注意：返回的Row字段顺序必须与MyItem定义的顺序一致（code在前，category在后），否则会出现类型不匹配问题。

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