Pyspark partition by most count

user_13312199 0
Viewed 865

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

pyspark window
1 Answers

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