Pyspark Cumcount with String Values

user_12231431 0
Viewed 104

My input dataframe is;

ID   Amount   result
45   1        default
46   2        default
47   3        default
48   5        default
48   10       default
49   12       default
50   5        default
50   7        default
50   9        default

Desired output is;

ID   Amount   result
45   1        default1
46   2        default1
47   3        default1
48   5        default1
48   10       default2
49   12       default1
50   5        default1
50   7        default2
50   9        default3

If ID is not changed default string should increase like this "default1", "default2". First is should always be "default1".

For python, i can do that with this code;

df['result'] += df.groupby('ID').cumcount().add(1).astype(str)

Could you please help me about for pyspark code?

pyspark apache-spark
1 Answers

Yes, you can use Spark window functions.

PS: If you don't have a column in your dataframe that could serve as an index to sort by, use the "ID" column itself. I've used an "index" column for the purpose of the example.

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window
from pyspark.sql.functions import concat, col, lit

df = spark.createDataFrame(
    [
        (1, 45, 1, 'default'),
        (2, 46, 2, 'default'),
        (3, 47, 3, 'default'),
        (4, 48, 5, 'default'),
        (5, 48, 10, 'default'),
        (6, 49, 12, 'default'),
        (7, 50, 5, 'default'),
        (8, 50, 7, 'default'),
        (9, 50, 9, 'default')
    ],
    ['index', 'ID', 'Amount', 'result'] # add your columns label here
)

df = df.withColumn("occurence", row_number().over(Window.partitionBy("ID").orderBy("index")))
df = df.withColumn("result", concat(col("result"), col("occurence"))).drop("occurence")
df.orderBy("index").show()

Result:

+-----+---+------+--------+
|index| ID|Amount|  result|
+-----+---+------+--------+
|    1| 45|     1|default1|
|    2| 46|     2|default1|
|    3| 47|     3|default1|
|    4| 48|     5|default1|
|    5| 48|    10|default2|
|    6| 49|    12|default1|
|    7| 50|     5|default1|
|    8| 50|     7|default2|
|    9| 50|     9|default3|
+-----+---+------+--------+
user_70957510
Related