Spark DataFrame 添加自增id

Spark DataFrame 添加自增id

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

方式一:利用窗口函数
// 加载数据
val dataframe = spark.read.csv(inputFile).toDF("lon", "lat")
    /**
      * 设置窗口函数的分区以及排序,因为是全局排序而不是分组排序,所有分区依据为空
      * 排序规则没有特殊要求也可以随意填写
      */
    val spec = Window.partitionBy().orderBy($"lon")

    val df1 = dataframe.withColumn("id", row_number().over(spec))

    df1.show()

+-----------+-----------+---+
|        lon|        lat| id|
+-----------+-----------+---+
|106.4273071|29.63554591|  1|
|  106.44104|29.51372023|  2|
|106.4602661|29.60211821|  3|
|106.4657593|29.45394812|  4|
+-----------+-----------+---+

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

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

方法二 利用RDD的 zipWithIndex算子

来自官网的描述
zipWithIndex算子

   // 在原Schema信息的基础上添加一列 “id”信息
    val schema: StructType = dataframe.schema.add(StructField("id", LongType))

    // DataFrame转RDD 然后调用 zipWithIndex
    val dfRDD: RDD[(Row, Long)] = dataframe.rdd.zipWithIndex()

    val rowRDD: RDD[Row] = dfRDD.map(tp => Row.merge(tp._1, Row(tp._2)))

    // 将添加了索引的RDD 转化为DataFrame
    val df2 = spark.createDataFrame(rowRDD, schema)

    df2.show()
+-----------+-----------+---+
|        lon|        lat| id|
+-----------+-----------+---+
|106.4273071|29.63554591|  0|
|  106.44104|29.51372023|  1|
|106.4602661|29.60211821|  2|
|106.4657593|29.45394812|  3|
+-----------+-----------+---+

第二种方式高效也简单。

参考资料

https://blog.csdn.net/u013090676/article/details/80379371

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值