CODEWITHSUNDEEP
apache-sparkapache-spark-sql
Affenbrotbaum
Affenbrotbaum

Reputation: 21

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

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.

Upvotes: 2

Views: 578

Answers (2)

Lamanus
Lamanus

Reputation: 13551

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

Upvotes: 0

kavetiraviteja
kavetiraviteja

Reputation: 2208

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.

Upvotes: 1

Related Questions