Spark -- 对DataFrame增加一列索引列(自增id列)==》(解决出现ID自增且唯一,但是不呈现自然数递增的问题)

时间:2024-03-31 11:13:53

Spark DataFrame 添加自增id
  在用Spark 处理数据的时候,经常需要给全量数据增加一列自增ID序号,在存入数据库的时候,自增ID也常常是一个很关键的要素。在使用mmlspark的LightGBMRanker时也需要指定一列int/long类型的id列,下面是几种实现方式。

方式一:利用RDD的 zipWithIndex算子,官网介绍如下

Spark -- 对DataFrame增加一列索引列(自增id列)==》(解决出现ID自增且唯一,但是不呈现自然数递增的问题)
// 在原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|
+-----------+-----------+---+


方式二:使用SparkSQL的function
    import org.apache.spark.sql.functions._
    val inputDF = inputDF.withColumn("id", monotonically_increasing_id)
    inputDF.show

 

源码介绍如下:

Spark -- 对DataFrame增加一列索引列(自增id列)==》(解决出现ID自增且唯一,但是不呈现自然数递增的问题)

注意:

monotonically_increasing_id这个方法,生成单调递增的64位整数的列表达式。

*

*生成的ID被保证是单调递增且唯一的,但不是连续的。

*当前的实现将分区ID放在高位31位,并将记录号

*在下33位的每个分区内。假设数据帧

*不到10亿个分区,每个分区有不到80亿条记录。

*例如,假设一个“DataFrame”有两个分区,每个分区有3条记录。

*此表达式将返回以下ID:

* {{{

*0,1,2,8589934592(1L<33),8589934593,8589934594。

* }}}

 

 

方式三:利用窗口函数

// 加载数据
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()