如何向n（PySpark）开始的火花数据帧添加顺序递增的列？

Question

我有一个数据框，其值如

 -------------
| col1 | col2 |       
 -------------
| a    |   2  |
| b    |   3  |
| c    |   4  |
| d    |   5  |
 -------------

我想创建第三列，该列具有从指定数字开始的自动递增值。

假设我希望数字从5开始，则预期输出应为

 --------------------
| col1 | col2 | col3 | 
|--------------------|
| a    |   2  | 5    |
| b    |   3  | 6    |
| c    |   4  | 7    |
| d    |   5  | 8    |
 --------------------

我已经尝试过monotonically_increasing_id()，但是由于分区，它没有给出连续的数字，并且它也没有以指定数字开头的功能。

我正在使用Spark 2.4.5和Pyspark（python）。

Answer 1

我们可以将窗口函数row_number与orderby()（用于全局排序）子句一起使用，并添加4以从5开始序列

在Pyspark中：

Using row_number():

from pyspark.sql.window import Window
from pyspark.sql.functions import *

df=spark.createDataFrame([('a',2),('b',3),('c',4),('d',5)],['col1','col2']).repartition(4)
w=Window.orderBy('col2')
df.withColumn("col3",row_number().over(w) + 4).show()
#+----+----+----+
#|col1|col2|col3|
#+----+----+----+
#|   a|   2|   5|
#|   b|   3|   6|
#|   c|   4|   7|
#|   d|   5|   8|
#+----+----+----+

在Scala中：

val w = Window.orderBy("col2")
df.withColumn("col3", row_number().over(w) + 4).show()

//+----+----+----+
//|col1|col2|col3|
//+----+----+----+
//|   a|   2|   5|
//|   b|   3|   6|
//|   c|   4|   7|
//|   d|   5|   8|
//+----+----+----+