from pyspark.sql import SparkSession
from pyspark.sql import Window, functions as F

spark = SparkSession.builder.getOrCreate()

data = [
('12a4', 'summer'),
('12a4', 'goest'),
('3b', 'yes'),
('3b', 'No'),
('3b', 'why'),
('4t', 'Hi')
]
df1 = spark.createDataFrame(data, ['id', 'col1'])
df1.show()

# +----+------+

# |  id|  col1|

# +----+------+

# |12a4|summer|

# |12a4| goest|

# |  3b|   yes|

# |  3b|    No|

# |  3b|   why|

# |  4t|    Hi|

# +----+------+

df = df1.select('id').distinct()
df = df.withColumn('new_id', F.row_number().over(Window.orderBy('id')))
df.show()

# +----+------+

# |  id|new_id|

# +----+------+

# |12a4|     1|

# |  3b|     2|

# |  4t|     3|

# +----+------+

df = df.join(df1, 'id', 'full')
df.show()

# +----+------+------+

# |  id|new_id|  col1|

# +----+------+------+

# |12a4|     1|summer|

# |12a4|     1| goest|

# |  4t|     3|    Hi|

# |  3b|     2|   yes|

# |  3b|     2|    No|

# |  3b|     2|   why|

# +----+------+------+

df = df.drop('id').withColumnRenamed('new_id', 'id')
df.show()

# +---+------+

# | id|  col1|

# +---+------+

# |  1|summer|

# |  1| goest|

# |  3|    Hi|

# |  2|   yes|

# |  2|    No|

# |  2|   why|

# +---+------+

df = df.filter(F.col('id') == 2)
df.show()

# +---+----+

# | id|col1|

# +---+----+

# |  2| yes|

# |  2|  No|

# |  2| why|

# +---+----+