Scala Spark Incrementing a column based on another column in dataframe without for loops

Question

I have a dataframe like the one below. I want a new column called cutofftype - which instead of the current monotonically increasing number should reset to 1 every time the ID column changes .

df = df.orderBy("ID","date").withColumn("cutofftype",monotonically_increasing_id()+1)

+------+---------------+----------+
|   ID |    date       |cutofftype|
+------+---------------+----------+
| 54441|     2016-06-20|         1|
| 54441|     2016-06-27|         2|
| 54441|     2016-07-04|         3|
| 54441|     2016-07-11|         4|
| 54500|     2016-05-02|         5|
| 54500|     2016-05-09|         6|
| 54500|     2016-05-16|         7|
| 54500|     2016-05-23|         8|
| 54500|     2016-06-06|         9|
| 54500|     2016-06-13|        10|
+------+---------------+----------+

Target is this as below :

+------+---------------+----------+
|   ID |    date       |cutofftype|
+------+---------------+----------+
| 54441|     2016-06-20|         1|
| 54441|     2016-06-27|         2|
| 54441|     2016-07-04|         3|
| 54441|     2016-07-11|         4|
| 54500|     2016-05-02|         1|
| 54500|     2016-05-09|         2|
| 54500|     2016-05-16|         3|
| 54500|     2016-05-23|         4|
| 54500|     2016-06-06|         5|
| 54500|     2016-06-13|         6|
+------+---------------+----------+

I know this can be done with for loops - i want to do it without for loops >> Is there a way out ?

Daeho Ro · Accepted Answer · 2020-09-08 10:57:15Z

1

Simple partition by problem. You should use the window.

import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy("ID").orderBy("date")

df.withColumn("cutofftype", row_number().over(w)).show()


+-----+----------+----------+
|   ID|      date|cutofftype|
+-----+----------+----------+
|54500|2016-05-02|         1|
|54500|2016-05-09|         2|
|54500|2016-05-16|         3|
|54500|2016-05-23|         4|
|54500|2016-06-06|         5|
|54500|2016-06-13|         6|
|54441|2016-06-20|         1|
|54441|2016-06-27|         2|
|54441|2016-07-04|         3|
|54441|2016-07-11|         4|
+-----+----------+----------+

answered Sep 8, 2020 at 10:57

Daeho Ro

13.7k4 gold badges25 silver badges50 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Scala Spark Incrementing a column based on another column in dataframe without for loops

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related