python - 如何将顺序递增的列添加到从 n 开始的 spark 数据帧（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)。

标签： pythonapache-sparkpysparkauto-increment

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|
#+----+----+----+

在斯卡拉：

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|
//+----+----+----+