Как добавить последовательно увеличивающийся столбец к искровому фрейму данных, который начинается с 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|
#+----+----+----+

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