Spark递归补全嵌套复杂类型DataFrame缺失列及代码问题排查

我有一个包含简单类型和复杂结构（如struct、struct数组、struct的数组的数组）的DataFrame，同时有一个以StructType为根的预期Schema。需要将该DataFrame调整为符合预期Schema的结构，若DataFrame中存在预期Schema里有但自身缺失的列，需添加这些列并设置默认值。

示例 #

预期Schema :-

root
 - struct
   - a: String
   - b: Int
   - c: array of struct
     - e: String
     - f: String

DataFrame Schema :-

root
 - struct
   - a: String
   - b: Int
   - c: array of String

如何使用Spark对每个层级元素进行递归处理？

我的代码 #

// Function to add missing columns recursively
def addMissingColumns(df1: DataFrame, df2: DataFrame, currentPath: String = ""): DataFrame = {
    val df1Schema = df1.schema
    val df2Schema = df2.schema

    val missingColumns = df1Schema.fields.filterNot { field1 =>
        df2Schema.fields.exists { field2 =>
            field1.name == field2.name && field1.dataType == field2.dataType
        }
    }

    val dfWithMissingColumns = missingColumns.foldLeft(df2) { (accDF, field) =>
        val colName = currentPath + field.name
        val dataType = field.dataType

        dataType match {
            case _: StructType =>
                val updatedDF = if (accDF.columns.contains(colName)) {
                    val missingCols = addMissingColumns(df1.selectExpr(s"$colName.*"), accDF.selectExpr(s"$colName.*"), colName + ".")
                    missingCols
                } else {
                    accDF.withColumn(colName, lit(null).cast(dataType)) // You can cast to the appropriate data type
                }
                updatedDF
            case _ =>
                // Handle non-struct types by adding them with null values if missing
                val updatedDF = if (accDF.columns.contains(colName)) {
                    accDF
                } else {
                    accDF.withColumn(field.name, lit(null).cast(dataType))
                }
                updatedDF
        }
    }
    dfWithMissingColumns.printSchema()
    dfWithMissingColumns.show()
    dfWithMissingColumns
}


def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
        .appName("YourAppName")
        .master("local[1]") // Use all available CPU cores
        .getOrCreate()

    // Sample DataFrames
    val df1 = spark.createDataFrame(Seq(
        (1, "Alice", Array(1, 2), (10, "New York", ("dddd","ggg"))),
        (2, "Bob", Array(3, 4), (20, "San Francisco", ("ddd","ggg")))
    )).toDF("id", "name", "numbers", "location")
    df1.printSchema()

    val df2 = spark.createDataFrame(Seq(
        (1, "Alice", Array(1, 2), (10, "New York 2")),
        (2, "Bob", Array(3, 4), (20, "San Francisco 2")),
        (3, "Ankur", Array(3, 4), (20, "Netherlands 2"))
    )).toDF("id", "name", "numbers", "location")

    df2.printSchema()

    addMissingColumns(df1, df2, "").printSchema()
}

当前问题 #

执行后仅返回location列相关内容，输出Schema如下：

root
 |-- _1: integer (nullable = true)
 |-- _2: string (nullable = true)
 |-- location._3: struct (nullable = true)
 |    |-- _1: string (nullable = true)
 |    |-- _2: string (nullable = true)

解决方案 #

问题分析 #

原代码的核心问题：

1. 处理嵌套Struct时，调用selectExpr(s"$colName.*")会丢弃原DataFrame的其他列，只返回struct展开后的字段
2. 补全struct内部缺失列后，没有将修改后的字段重新嵌套为原struct列，而是直接返回展开的字段
3. 路径拼接和列名处理逻辑错误，导致非顶层列的命名混乱

修改后的递归实现 #

import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{ArrayType, StructType}

object SchemaAlignment {
  // 递归生成符合目标Schema的列表达式
  private def alignColumn(targetField: org.apache.spark.sql.types.StructField, sourceDF: DataFrame): org.apache.spark.sql.Column = {
    val colName = targetField.name
    val targetType = targetField.dataType

    // 检查源DataFrame是否存在该列
    val sourceColExists = sourceDF.columns.contains(colName)

    targetType match {
      case structType: StructType =>
        // 处理嵌套Struct：递归生成每个子列的表达式，再组装成Struct
        val structFields = structType.fields.map { subField =>
          alignColumn(subField, if (sourceColExists) sourceDF.select(s"$colName.*") else sourceDF).alias(subField.name)
        }
        if (sourceColExists) {
          struct(structFields: _*).alias(colName)
        } else {
          // 源列不存在时，创建空Struct
          struct(structFields: _*).cast(structType).alias(colName)
        }

      case arrayType: ArrayType if arrayType.elementType.isInstanceOf[StructType] =>
        // 处理数组嵌套Struct：先处理数组内的Struct元素，再包装成数组
        val elementStructType = arrayType.elementType.asInstanceOf[StructType]
        val elementExpr = if (sourceColExists) {
          transform(col(colName), elem => {
            val subFields = elementStructType.fields.map(subField => {
              val elemCol = elem.getField(subField.name)
              if (elemCol.isNullAt(0) && !sourceDF.select(col(colName).getField(subField.name)).isEmpty) {
                lit(null).cast(subField.dataType).alias(subField.name)
              } else {
                alignColumn(subField, sourceDF.select(col(colName).getField(subField.name))).alias(subField.name)
              }
            })
            struct(subFields: _*)
          })
        } else {
          // 源列不存在时，创建空数组
          array().cast(arrayType).alias(colName)
        }
        elementExpr.cast(arrayType).alias(colName)

      case _ =>
        // 处理简单类型或普通数组：存在则用源列，不存在则设为默认值（null）
        if (sourceColExists) {
          col(colName).cast(targetType).alias(colName)
        } else {
          lit(null).cast(targetType).alias(colName)
        }
    }
  }

  // 对外暴露的对齐Schema方法
  def alignSchema(sourceDF: DataFrame, targetSchema: StructType): DataFrame = {
    val alignedColumns = targetSchema.fields.map(field => alignColumn(field, sourceDF))
    sourceDF.select(alignedColumns: _*)
  }

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("SchemaAlignment")
      .master("local[*]")
      .getOrCreate()
    import spark.implicits._

    // 示例1：测试嵌套Struct补全
    val targetSchema1 = StructType(Seq(
      org.apache.spark.sql.types.StructField("id", org.apache.spark.sql.types.IntegerType),
      org.apache.spark.sql.types.StructField("name", org.apache.spark.sql.types.StringType),
      org.apache.spark.sql.types.StructField("numbers", org.apache.spark.sql.types.ArrayType(org.apache.spark.sql.types.IntegerType)),
      org.apache.spark.sql.types.StructField("location", StructType(Seq(
        org.apache.spark.sql.types.StructField("zip", org.apache.spark.sql.types.IntegerType),
        org.apache.spark.sql.types.StructField("city", org.apache.spark.sql.types.StringType),
        org.apache.spark.sql.types.StructField("detail", StructType(Seq(
          org.apache.spark.sql.types.StructField("street", org.apache.spark.sql.types.StringType),
          org.apache.spark.sql.types.StructField("district", org.apache.spark.sql.types.StringType)
        )))
      )))
    ))

    val sourceDF1 = spark.createDataFrame(Seq(
      (1, "Alice", Array(1, 2), (10, "New York 2")),
      (2, "Bob", Array(3, 4), (20, "San Francisco 2")),
      (3, "Ankur", Array(3, 4), (20, "Netherlands 2"))
    )).toDF("id", "name", "numbers", "location")

    println("源DataFrame Schema:")
    sourceDF1.printSchema()

    val alignedDF1 = alignSchema(sourceDF1, targetSchema1)
    println("\n对齐后的Schema:")
    alignedDF1.printSchema()
    println("\n对齐后的数据:")
    alignedDF1.show(false)

    // 示例2：测试数组嵌套Struct的转换
    val targetSchema2 = StructType(Seq(
      org.apache.spark.sql.types.StructField("a", org.apache.spark.sql.types.StringType),
      org.apache.spark.sql.types.StructField("b", org.apache.spark.sql.types.IntegerType),
      org.apache.spark.sql.types.StructField("c", org.apache.spark.sql.types.ArrayType(StructType(Seq(
        org.apache.spark.sql.types.StructField("e", org.apache.spark.sql.types.StringType),
        org.apache.spark.sql.types.StructField("f", org.apache.spark.sql.types.StringType)
      ))))
    ))

    val sourceDF2 = spark.createDataFrame(Seq(
      ("test1", 1, Array("val1", "val2")),
      ("test2", 2, Array("val3", "val4"))
    )).toDF("a", "b", "c")

    println("\n\n源DataFrame Schema（数组为字符串类型）:")
    sourceDF2.printSchema()

    val alignedDF2 = alignSchema(sourceDF2, targetSchema2)
    println("\n对齐后的Schema（数组转为Struct类型）:")
    alignedDF2.printSchema()
    println("\n对齐后的数据:")
    alignedDF2.show(false)

    spark.stop()
  }
}

关键改进点 #

1. 全列保留：通过select(alignedColumns: _*)方式构建新DataFrame，确保所有目标Schema的列都被保留，不会丢失原列
2. 递归嵌套处理：对Struct类型递归生成子列表达式，再用struct()重新组装为原列；对数组嵌套Struct，用transform()函数处理每个数组元素
3. 类型兼容转换：自动将源列类型转换为目标类型，不存在的列设置为null（可根据需求修改默认值）
4. 路径正确处理：避免了原代码中路径拼接错误的问题，直接通过字段名递归处理嵌套结构

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