Pyspark partition by most count - pyspark

Question

I have a question,
I have a dataset like this :
id / color
1 / red
2 / green
2 / green
2 / blue
3 / blue
4 / yellow
4 / pink
5 / red
and I would like to group by id and to keep the most frequent color
to have something like this :
(in case of draw take randomly it's ok, or better solution if you have)
id / most_color
1 / red
2 / green
3 / blue
4 / yellow
5 / red
I have tried things like :
display(dataset.select("id","color").
dropDuplicates().
withColumn("most_color",count("color").over(w)))
or like this :
dataset2= (dataset.select("id","color").
withColumn("most_color", dataset["color"]).
groupBy("id").
agg(count('color').
alias('count').
filter(column('count') == max(count))))
display(dataset2)
thank you everyone

Answer 1

You can use Window function row_number() to achieve this
from pyspark.sql import functions as F
from pyspark.sql import Window as W
_w = W.partitionBy('id').orderBy(F.col('id').desc())
_w = W.partitionBy('id').orderBy(F.col('id').desc())
df_final = df_final.withColumn('rn_no', F.row_number().over(_w))
df_final = df_final.filter(F.col('rn_no') == 1)
df_final.show()
Output
id / most_color
1 / red
2 / green
3 / blue
4 / yellow
5 / red
Modified Version : This will give you the most used/appeared value in a group --
Input
df_a = spark.createDataFrame([(1,'red'),(2,'green'),(2,'green'),(2,'blue'),(3,'blue'),(4,'yellow'),(4,'pink'),(5,'red')],[ "id","color"])
+---+------+
| id| color|
+---+------+
| 1| red|
| 2| green|
| 2| green|
| 2| blue|
| 3| blue|
| 4|yellow|
| 4| pink|
| 5| red|
+---+------+
# First Group the values to get the max appeared color in a group
df = df_a.groupBy('id','color').agg(F.count('color').alias('count')).orderBy(F.col('id'))
# Now, make a partition and sort of the decending order for each window of ID and take the first value
_w = W.partitionBy('id').orderBy(F.col('count').desc())
df_a = df.withColumn('rn_no', F.row_number().over(_w))
df_a = df_a.filter(F.col('rn_no') == F.lit('1'))
Output
df_a.show()
+---+-----+-----+-----+
| id|color|count|rn_no|
+---+-----+-----+-----+
| 1| red| 1| 1|
| 2|green| 2| 1|
| 3| blue| 1| 1|
| 4| pink| 1| 1|
| 5| red| 1| 1|
+---+-----+-----+-----+