Spark DataFrame导入mysql入库添加自增主键id

Spark DataFrame 添加一个自增主键id

  在使用SparkSQL处理数据的时候，经常需要给全量数据增加一列自增的ID序号，在存入数据库的时候，自增ID也常常是一个很关键的要素。在DataFrame的API中没有实现这一功能，所以只能通过其他方式实现，或者转成RDD再用RDD的zipWithIndex算子实现。
下面介绍两种实现方式：

方式一： 利用窗口函数

// 加载数据
val df: DataFrame = session.read.option("header", "true").csv(inPath)

/**
      * 设置窗口函数的分区及排序，因为是全局排序而不是分组排序，所有分区依据为空
      * 排序规则没有特殊要求也可以随意填写
      */
    val spec: WindowSpec = Window.partitionBy().orderBy(col("imo"))
    val result: DataFrame = df.withColumn("id", row_number().over(spec))
    result.show()

该方法的缺点就是在添加自增ID的过程中，会将所有的数据集中到一个分区，会导致性能下降。

WARN [main] - No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.

方法二： 利用RDD的zipWithIndex算子
来自官网的描述

// 在原Schema信息的基础上添加一列“id”信息
    val schema: StructType = res.schema.add(StructField("id", LongType))
    
    // DataFrame转RDD, 然后调用 zipWithIndex
    val dfRDD: RDD[(Row, Long)] = res.rdd.zipWithIndex()
    
    // 将id字段合并在一起，merge顺序不可修改，因为添加id的schema字段在最后一个
    val rowRDD: RDD[Row] = dfRDD.map(tp => Row.merge(tp._1,Row(tp._2)))
    
    // 将添加了索引的RDD转化为 DataFrame
    val result: DataFrame = session.createDataFrame(rowRDD, schema)

将结果进行保存到mysql：

val prop =new Properties()
prop.setProperty("user","root")
prop.setProperty("password","root")
result.write.mode(SaveMode.Append).jdbc("jdbc:mysql://localhost:3306/loong","static_recently",prop)

第二种方式高效也简单