Spark 2.X及以下版本如何为嵌套数组结构体列新增字段

问题场景 #

现有DataFrame的Schema如下：

root
     |-- date: timestamp (nullable = true)
     |-- questionAnswerList: array (nullable = true)
     |    |-- element: struct (containsNull = true)
     |    |    |-- questionNumber: string (nullable = true)
     |    |    |-- listAnswers: array (nullable = true) 
     |    |    |    |-- element: string(containsNull = true)

需要为questionAnswerList数组内的每个结构体新增index整数字段，目标Schema如下：

root
     |-- date: timestamp (nullable = true)
     |-- questionAnswerList: array (nullable = true)
     |    |-- element: struct (containsNull = true)
     |    |    |-- index: integer (nullable = true)
     |    |    |-- questionNumber: string (nullable = true)
     |    |    |-- listAnswers: array (nullable = true) 
     |    |    |    |-- element: string(containsNull = true)

用户尝试了以下UDF但报错：

val  addIndexInStruct: UserDefinedFunction = udf((data: Seq[Row]) => {
    data.zipWithIndex.map{case (Row(x:String,y:Array[String]), index) => (index, x, y )}
  })

df.withColumn("newCol",addIndexInStruct($"questionAnswerList")).show(false)

错误信息：

Caused by: scala.MatchError: ([Q10,WrappedArray(R10.1, R10.2)],0) (of class scala.Tuple2)

错误原因 #

zipWithIndex会将数组的每个元素与对应的索引组成Tuple2(元素, 索引)，而你的模式匹配写成了case (Row(...), index)，相当于把整个Tuple2当成了第一个参数，自然匹配失败。此外，直接返回Tuple无法对应Struct类型，需要构造符合结构的Row对象。

解决方案 #

方法1：修正UDF实现 #

调整模式匹配的顺序，并构造正确的Struct类型返回：

import org.apache.spark.sql.Row
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types.{ArrayType, IntegerType, StringType, StructField, StructType}

// 定义目标结构体的Schema
val targetStructType = StructType(
  List(
    StructField("index", IntegerType, nullable = true),
    StructField("questionNumber", StringType, nullable = true),
    StructField("listAnswers", ArrayType(StringType), nullable = true)
  )
)

val addIndexInStruct = udf((data: Seq[Row]) => {
  data.zipWithIndex.map { case (row, index) =>
    // 从原Row中提取字段，再新增index字段构造新Row
    Row(index, row.getAs[String]("questionNumber"), row.getAs[Seq[String]]("listAnswers"))
  }
}, ArrayType(targetStructType))

// 替换原列或新增列
val resultDF = df.withColumn("questionAnswerList", addIndexInStruct($"questionAnswerList"))
resultDF.printSchema()

方法2：无需UDF的实现（推荐，避免UDF性能开销） #

利用explode展开数组，结合窗口函数生成索引，再重新聚合：

import org.apache.spark.sql.functions.{col, explode, row_number}
import org.apache.spark.sql.expressions.Window

// 1. 展开数组，为每个结构体生成索引
val explodedDF = df.withColumn("struct_element", explode($"questionAnswerList"))
  .withColumn("index", row_number().over(Window.partitionBy("date").orderBy("struct_element.questionNumber")) - 1) // 索引从0开始

// 2. 构造带index的新结构体
val newStructDF = explodedDF.withColumn(
  "new_struct",
  struct(
    col("index"),
    col("struct_element.questionNumber"),
    col("struct_element.listAnswers")
  )
)

// 3. 重新聚合数组，恢复原结构
val resultDF = newStructDF.groupBy("date")
  .agg(collect_list("new_struct").alias("questionAnswerList"))

resultDF.printSchema()

注意：如果date不是唯一标识行的字段，需要替换为能唯一区分每一行的主键或组合字段，避免聚合错误。

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