如何用Scala高效构建符合指定Schema的目标Dataset？

[c1, c3, c6] [aaa, [ab, ac], [bc, bd]]
[c2, c3, c4, c5] [bbb, [bb, bc], ccc, ddd]
[c4, c6, c7] [fff, [ss, pp, l], zzz]

### 期望输出Dataset

c1 c2 c3 c4 c5 c6 c7 #

aaa null [ab, ac] null null [bc, bd] null
null bbb [bb, bc] ccc ddd null null
null null null fff null [ss, pp, l] zzz

请问如何用Scala高效创建符合该Schema的Dataset？

---

# 解决方案
这个需求很好实现，核心思路就是先明确目标Schema，再把每行的键值对数组映射成符合Schema的Row对象，最后生成新的Dataset。下面是具体步骤和代码：

### 1. 定义目标Schema
首先要确定最终输出的字段类型，按照需求，`c3`和`c6`是字符串数组类型，其余字段为字符串类型：
```scala
import org.apache.spark.sql.types._

val targetSchema = StructType(Seq(
  StructField("c1", StringType, nullable = true),
  StructField("c2", StringType, nullable = true),
  StructField("c3", ArrayType(StringType), nullable = true),
  StructField("c4", StringType, nullable = true),
  StructField("c5", StringType, nullable = true),
  StructField("c6", ArrayType(StringType), nullable = true),
  StructField("c7", StringType, nullable = true)
))

2. 编写行转换辅助函数 #

接下来需要把原Dataset的每一行（col1和col2的数组对）转换成Row对象——不存在的字段填null，存在的字段要匹配对应类型：

import org.apache.spark.sql.Row

def convertToRow(col1: Seq[String], col2: Seq[Any]): Row = {
  // 将col1和col2打包成键值对Map，方便快速查找
  val keyValueMap = (col1 zip col2).toMap
  
  // 按照targetSchema的字段顺序逐个取值，处理类型匹配
  val values = targetSchema.fields.map { field =>
    keyValueMap.get(field.name) match {
      case Some(value) => 
        // 数组类型字段转换为Seq[String]，对应Spark的ArrayType
        if (field.dataType.isInstanceOf[ArrayType]) value.asInstanceOf[Seq[String]]
        else value.asInstanceOf[String]
      case None => null
    }
  }
  
  Row.fromSeq(values)
}

3. 应用转换生成最终Dataset #

假设你的原Dataset叫sourceDf，我们可以用RDD API做转换（无Shuffle开销，效率很高），然后生成新的Dataset：

import org.apache.spark.sql.SparkSession

val spark = SparkSession.builder().getOrCreate()

// 将原DataFrame转为RDD，逐行应用转换函数
val resultRDD = sourceDf.rdd.map { row =>
  val col1 = row.getAs[Seq[String]]("col1")
  val col2 = row.getAs[Seq[Any]]("col2")
  convertToRow(col1, col2)
}

// 用目标Schema创建最终的DataFrame
val resultDf = spark.createDataFrame(resultRDD, targetSchema)

// 查看结果
resultDf.show(false)

4. 可选：用UDF替代RDD转换 #

如果不想切换到RDD API，也可以用UDF实现，代码如下：

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

val convertUdf = udf((col1: Seq[String], col2: Seq[Any]) => {
  val keyValueMap = (col1 zip col2).toMap
  targetSchema.fields.map { field =>
    keyValueMap.get(field.name) match {
      case Some(v) => 
        if (field.dataType.isInstanceOf[ArrayType]) v.asInstanceOf[Seq[String]] 
        else v.asInstanceOf[String]
      case None => null
    }
  }
})

// 将UDF的结果拆分为各个字段
val resultDf = sourceDf.withColumn("temp", convertUdf($"col1", $"col2"))
  .select(targetSchema.fields.zipWithIndex.map { case (field, idx) =>
    $"temp".getItem(idx).cast(field.dataType).alias(field.name)
  }:_*)

小提醒 #

• 要确保原col2中的数组值确实是Seq[String]类型，如果是其他类型（比如Array），可以提前转换为Seq
• 这两种方法都是逐行操作，没有额外的Shuffle开销，处理大数据量时效率优异

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