We are trying to get 'Expected Output' start from the highest number from 'Doc_number' grouped by 'Roll_NO'.
| Roll_NO | Doc_number | Expected Output |
|---|---|---|
| 171 | 2 | 8 |
| 171 | 3 | 9 |
| 171 | 4 | 10 |
| 171 | 5 | 11 |
| 171 | 6 | 12 |
| 171 | 7 | 13 |
| 181 | 2 | 7 |
| 181 | 3 | 8 |
| 181 | 4 | 9 |
| 181 | 5 | 10 |
| 181 | 6 | 11 |
| 183 | 2 | 5 |
| 183 | 3 | 6 |
| 183 | 4 | 7 |
How to pass max to to the next column and start from that number?
You can do it using window functions.
Initial df:
df = spark.createDataFrame(
[(171, 2),
(171, 3),
(171, 4),
(171, 5),
(171, 6),
(171, 7),
(181, 2),
(181, 3),
(181, 4),
(181, 5),
(181, 6),
(183, 2),
(183, 3),
(183, 4)],
['Roll_NO', 'Doc_number']
)
Script:
from pyspark.sql import functions as F, Window as W
w_desc = W.partitionBy('Roll_NO').orderBy(F.desc('Doc_number'))
w_asc = W.partitionBy('Roll_NO').orderBy('Doc_number')
df = df.withColumn('output', F.max('Doc_number').over(w_desc) + F.row_number().over(w_asc))
df.show()
# +-------+----------+------+
# |Roll_NO|Doc_number|output|
# +-------+----------+------+
# | 171| 2| 8|
# | 171| 3| 9|
# | 171| 4| 10|
# | 171| 5| 11|
# | 171| 6| 12|
# | 171| 7| 13|
# | 181| 2| 7|
# | 181| 3| 8|
# | 181| 4| 9|
# | 181| 5| 10|
# | 181| 6| 11|
# | 183| 2| 5|
# | 183| 3| 6|
# | 183| 4| 7|
# +-------+----------+------+