如何基于已有Spark DataFrame生成符合指定结构的新DataFrame

你原有实现的核心问题如下：

• 调用collect()将所有数据拉取到Driver节点内存，数据量稍大就会触发OOM，完全没有利用Spark分布式计算的特性
• 通过索引下标取Row的字段，CSV列顺序一旦变动就会拿错值，可维护性极差
• 对split操作没有做边界校验，遇到source_table格式不符合schema.table规则、或者文件名不符合schema.table.csv规则的情况，直接报数组越界异常
• 手动实现JSON序列化属于重复造轮子，Spark原生支持直接将DataFrame导出为标准JSON格式文件，不需要额外序列化工具

1. 分布式转换逻辑（不将数据拉取到Driver） #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.DataFrame

def processDataLinkFile(file: LocatedFileStatus): DataFrame = {
    // 解析目标库表名，增加边界校验避免越界
    val fileNameParts = file.getPath.getName.split("\\.")
    require(fileNameParts.length >=3, s"文件名${file.getPath.getName}不符合schema.table.csv格式要求")
    val schemaTo = fileNameParts(0)
    val tableTo = fileNameParts(1)

    // 读取CSV文件
    val rawDf = spark.read
      .option("delimiter", ",")
      .option("header", true)
      .option("inferSchema", false)
      .csv(file.getPath.toUri.getPath)

    // 分布式转换生成符合要求的DataFrame
    rawDf
      // 拆分source_table为schema_from和table_from，过滤格式不符合的行
      .withColumn("source_table_parts", split(col("source_table"), "\\."))
      .filter(size(col("source_table_parts")) === 2)
      .withColumn("schema_from", col("source_table_parts").getItem(0))
      .withColumn("table_from", col("source_table_parts").getItem(1))
      // 映射其余字段
      .withColumn("column_from", col("source"))
      .withColumn("link_type", col("relation_type"))
      .withColumn("schema_to", lit(schemaTo))
      .withColumn("table_to", lit(tableTo))
      .withColumn("column_to", col("target"))
      // 选择对应顺序的字段，需要强类型的话可以直接.as[DataLink]转为DataSet
      .select(
          "schema_from",
          "table_from",
          "column_from",
          "link_type",
          "schema_to",
          "table_to",
          "column_to"
      )
}

2. 直接导出JSON文件 #

不需要手动实现序列化逻辑，直接用Spark原生的写JSON能力：

val resultDf = processDataLinkFile(yourFile)
// 小数据量需要单文件输出时加coalesce(1)，大数据量场景不要使用
resultDf
  .coalesce(1) 
  .write
  .option("encoding", "UTF-8")
  .mode("overwrite")
  .json("/path/to/your/output/directory")

可选：小数据量场景下直接获取JSON数组 #

如果确实需要把结果拉到Driver生成JSON数组字符串，直接用Spark内置的toJSON方法：

val jsonArray = "[" + resultDf.toJSON.collect().mkString(",") + "]"

• 全程使用Spark分布式算子计算，数据不会全部拉到Driver，支持大文件处理
• 通过列名取数，CSV列顺序变动不影响逻辑，可维护性更高
• 增加了文件名和source_table的格式校验，避免数组越界异常
• 直接使用Spark原生的JSON导出能力，不需要引入额外的JSON序列化依赖，性能更高，输出格式标准
• 如果需要强类型校验，直接调用.as[DataLink]即可转换为DataSet[DataLink]，实现类型安全

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