How to find multi mode of an array column in Pyspark

Viewed 491

I want to find the mode of the task column in this dataframe:

+-----+-----------------------------------------+
|  id |              task                       |
+-----+-----------------------------------------+
| 101 |   [person1, person1, person3]           |
| 102 |   [person1, person2, person3]           |
| 103 |           null                          |
| 104 |   [person1, person2]                    |
| 105 |   [person1, person1, person2, person2]  |
| 106 |           null                          |
+-----+-----------------------------------------+

If there are multiple modes, I want to display all of them.
Can someone please help me get this output:

+-----+-----------------------------------------+---------------------------+
|  id |              task                       |           mode            |
+-----+-----------------------------------------+---------------------------+
| 101 |   [person1, person1, person3]           |[person1]                  |
| 102 |   [person1, person2, person3]           |[person1, person2, person3]|
| 103 |           null                          |[]                         |
| 104 |   [person1, person2]                    |[person1, person2]         |
| 105 |   [person1, person1, person2, person2]  |[person1, person2]         |
| 106 |           null                          |[]                         |
+-----+-----------------------------------------+---------------------------+

This is my first question here. Any help or hint is greatly appreciated. Thank you.

3 Answers

I dont see a reason to use a UDF for this case for spark2.4+ as we can use higher order functions to get desired output. UDF using counter will be very slow for Big Data in comparison to higher order functions:

from pyspark.sql import functions as F

df\
  .withColumn("most_common", F.expr("""transform(array_distinct(values),\
                                      x-> array(aggregate(values, 0,(acc,t)->acc+IF(t=x,1,0)),x))"""))\
  .withColumn("most_common", F.expr("""transform(filter(most_common, x-> x[0]==array_max(most_common)[0]),y-> y[1])"""))\
  .show(truncate=False)

#+---+----------------------------------------+---------------------------+
#|id |values                                  |most_common                |
#+---+----------------------------------------+---------------------------+
#|1  |[good, good, good, bad, bad, good, good]|[good]                     |
#|2  |[bad, badd, good, bad,, good, bad, good]|[bad, good]                |
#|2  |[person1, person2, person3]             |[person1, person2, person3]|
#|2  |null                                    |null                       |
#+---+----------------------------------------+---------------------------+

Using Spark 2.3:

You can solve this using a custom UDF. For the purposes of getting multiple mode values, I'm using a Counter. I use the except block in the UDF for the null cases in your task column.
(For Python 3.8+ users, there is a statistics.multimode() in-built function you can make use of)

Your dataframe:

from pyspark.sql.types import *
from pyspark.sql import functions as F
from pyspark.sql.functions import *

schema = StructType([StructField("id", IntegerType()), StructField("task", ArrayType(StringType()))])
data = [[101, ["person1", "person1", "person3"]], [102, ["person1", "person2", "person3"]], [103, None], [104, ["person1", "person2"]], [105, ["person1", "person1", "person2", "person2"]], [106, None]]

df = spark.createDataFrame(data,schema=schema)

Operation:

from collections import Counter

def get_multi_mode_list(input_array):
    multi_mode = []
    counter_var = Counter(input_array)  
    try:
        temp = counter_var.most_common(1)[0][1]
    except:
        temp = counter_var.most_common(1)
    for i in counter_var: 
        if input_array.count(i) == temp: 
            multi_mode.append(i)
    return(list(set(multi_mode)))


get_multi_mode_list_udf = F.udf(get_multi_mode_list, ArrayType(StringType()))

df.withColumn("multi_mode", get_multi_mode_list_udf(col("task"))).show(truncate=False)

Output:

+---+------------------------------------+---------------------------+
|id |task                                |multi_mode                 |
+---+------------------------------------+---------------------------+
|101|[person1, person1, person3]         |[person1]                  |
|102|[person1, person2, person3]         |[person2, person3, person1]|
|103|null                                |[]                         |
|104|[person1, person2]                  |[person2, person1]         |
|105|[person1, person1, person2, person2]|[person2, person1]         |
|106|null                                |[]                         |
+---+------------------------------------+---------------------------+
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from collections import Counter
from itertools import groupby
from pyspark.sql.types import *
from datetime import datetime
from pyspark.sql import *
from collections import *
from pyspark.sql.functions import udf,explode
from pyspark.sql.types import StringType
table_schema = StructType([
                     StructField('key2', IntegerType(), True),
                 
                     StructField('list6', ArrayType(StringType()), False)
                     ])
df= spark.createDataFrame([
( 101 ,   ["person1", "person1", "person3"] ),
(102 ,   ["person1", "person2", "person3"]   ),
( 103 ,None  ),
( 104  ,  ["person1", "person2"]),
(105 ,   ["person1", "person1", "person2", "person2"])] 
,["id","List"])
def mode(list1):
    res = []
    if(list1 is None or len(list1)==0):
        return []
    test_list1 = Counter(list1)  
    temp = test_list1.most_common(1)[0][1]
    for ele in list1:
       if list1.count(ele) == temp:
           res.append(ele)
    return list(set(res))




df.createOrReplaceTempView("A")
spark.udf.register("mode", mode,ArrayType(StringType()))
spark.sql("select id,list,mode(list)func from A").show(truncate=False)
Related