Append a monotonically increasing id column that increases on column value match

user_14148985 0
Viewed 327

I am ingesting a dataframe and I want to append a monotonically increasing column that increases whenever another column matches a certain value. For example I have the following table

+------+-------+
| Col1 | Col2  |
+------+-------+
| B    |   543 |
| A    |  1231 |
| B    | 14234 |
| B    | 34234 |
| B    |  3434 |
| A    | 43242 |
| B    | 43242 |
| B    | 56453 |
+------+-------+

I would like to append a column that increases in value whenever "A" in col1 is present. So the result would look like

+------+-------+------+
| Col1 | Col2  | Col3 |
+------+-------+------+
| B    |   543 |    0 |
| A    |  1231 |    1 |
| B    | 14234 |    1 |
| B    | 34234 |    1 |
| B    |  3434 |    1 |
| A    | 43242 |    2 |
| B    | 43242 |    2 |
| B    | 56453 |    2 |
+------+-------+------+

Keeping the initial order is important.

I tried zippering but that doesn't seem to produce the right result. Splitting it up into individual seqs manually and doing it that way is not going to be performant enough (think 100+ GB tables). I looked into trying this with a map function that would keep a counter somewhere but couldn't get that to work.

Any advice or pointer in the right direction would be greatly appreciated.

apache-spark-sql apache-spark
2 Answers

spark does not provide any default functions to achieve this kind of functionality

I would do like to do most probably in this way

//inputDF contains Col1 | Col2 
val df = inputDF.select("Col1").distinct.rdd.zipWithIndex().toDF("Col1","Col2")
val finalDF = inputDF.join(df,df("Col1") === inputDF("Col1"),"left").select(inputDF("*"),"Col3")

but the problem here I can see is (join which will result in the shuffle).

you can also check other autoincrement API's here.

user_34754880

Use window and sum over the window of the value 1 when Col1 = A.

import pyspark.sql.functions as f
from pyspark.sql import Window

w = Window.partitionBy().rowsBetween(Window.unboundedPreceding, Window.currentRow)

df.withColumn('Col3', f.sum(f.when(f.col('Col1') == f.lit('A'), 1).otherwise(0)).over(w)).show()

+----+-----+----+
|Col1| Col2|Col3|
+----+-----+----+
|   B|  543|   0|
|   A| 1231|   1|
|   B|14234|   1|
|   B|34234|   1|
|   B| 3434|   1|
|   A|43242|   2|
|   B|43242|   2|
|   B|56453|   2|
+----+-----+----+
user_118415710
Related