How to reset the row number to 1 on a non consecutive record

user_6458707 0
Viewed 76

I have a hive table like the following. I want to assign the row number for user consecutive data. When the data is not consecutive, the row number will reset to 1.

| Date       | User |
|------------|------|
| 2022-05-01 | A    |
| 2022-05-02 | A    |
| 2022-05-03 | A    |
| 2022-05-05 | A    |
| 2022-05-06 | A    |
| 2022-05-01 | B    |
| 2022-05-03 | B    |
| 2022-05-04 | B    |
| 2022-05-05 | B    |
| 2022-05-06 | B    |
| 2022-05-01 | C    |
| 2022-05-02 | C    |
| 2022-05-04 | C    |
| 2022-05-05 | C    |
| 2022-05-06 | C    |

Here is the output I want.


| Date       | User | Rank |
|------------|------|------|
| 2022-05-01 | A    | 1    |
| 2022-05-02 | A    | 2    |
| 2022-05-03 | A    | 3    |
| 2022-05-05 | A    | 1    |
| 2022-05-06 | A    | 2    |
| 2022-05-01 | B    | 1    |
| 2022-05-03 | B    | 1    |
| 2022-05-04 | B    | 2    |
| 2022-05-05 | B    | 3    |
| 2022-05-06 | B    | 4    |
| 2022-05-01 | C    | 1    |
| 2022-05-02 | C    | 2    |
| 2022-05-04 | C    | 1    |
| 2022-05-05 | C    | 2    |
| 2022-05-06 | C    | 3    |

I couldn't figure out with row_number, lag and other udfs.

hiveql apache-spark-sql apache-spark hive sql
1 Answers

I was able to solve it with a combination of windows, lag and dense_rank, the key element to this solution is is the prefix sum over the window:

import spark.implicits._
val columns = Seq("Date", "User")
val data = Seq(("2022-05-01", "A"), ("2022-05-02", "A"), ("2022-05-03", "A"), ("2022-05-05", "A"), ("2022-05-06", "A"),
  ("2022-05-01", "B"), ("2022-05-03", "B"), ("2022-05-04", "B"), ("2022-05-05", "B"), ("2022-05-06", "B"),
  ("2022-05-01", "C"), ("2022-05-02", "C"), ("2022-05-04", "C"), ("2022-05-05", "C"), ("2022-05-06", "C")
)
val df = spark.sparkContext.parallelize(data).toDF(columns: _*)
val w = Window.partitionBy("User").orderBy("Date")
val w2 = Window.partitionBy("User", "partition").orderBy("Date")
val results = df.withColumn("old_date", lag("Date", 1, null).over(w))
  .withColumn("date_diff", datediff(col("date"),col("old_date")))
  .withColumn("tmp", when(col("old_date").isNull.or(col("date_diff").equalTo(1)), 0).otherwise(1))
  .withColumn("partition", sum("tmp").over(w.rowsBetween(Window.unboundedPreceding, Window.currentRow)))
  .withColumn("Rank", dense_rank().over(w2))
  .select("Date", "User", "Rank")
results.show()

Result:

+----------+----+----+
|      Date|User|Rank|
+----------+----+----+
|2022-05-01|   A|   1|
|2022-05-02|   A|   2|
|2022-05-03|   A|   3|
|2022-05-05|   A|   1|
|2022-05-06|   A|   2|
|2022-05-01|   B|   1|
|2022-05-03|   B|   1|
|2022-05-04|   B|   2|
|2022-05-05|   B|   3|
|2022-05-06|   B|   4|
|2022-05-01|   C|   1|
|2022-05-02|   C|   2|
|2022-05-04|   C|   1|
|2022-05-05|   C|   2|
|2022-05-06|   C|   3|
+----------+----+----+
user_63259940
Related