Spark DataFrame 添加自增id
在用Spark 处理数据的时候,经常需要给全量数据增加一列自增ID序号,在存入数据库的时候,自增ID也常常是一个很关键的要素。在使用mmlspark的LightGBMRanker时也需要指定一列int/long类型的id列,下面是几种实现方式。
利用RDD的 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|
+-----------+-----------+---+
使用SparkSQL的function
import org.apache.spark.sql.functions._
val inputDF = inputDF.withColumn("id", monotonically_increasing_id)
inputDF.show