find max value in list in a Spark dataframe

Viewed 717

I want to find the max value in a list. How do you do this in pyspark?

df = spark.createDataFrame([(1, [4,1]), (2, [4,5]), (3, [4,0])], ["A", "B"])
df.show()

+---+------+
|  A|     B|
+---+------+
|  1|[4, 1]|
|  2|[4, 5]|
|  3|[4, 0]|
+---+------+

In this following example, how do I find on each row the max value from the list in column B. So for:

  • row 1 --> 4
  • row 2 --> 5
  • row 3 --> 4
2 Answers

You can use aggregate function.

df = spark.createDataFrame([(1, [4,1, 4, 54,4, 2,2, 7,14, 23,74,53]), (2, [4,5, 11, 3,45, 34, 2,3, 4]), (3, [4,0, 32, 23,23, 5,23,2 ,37,8, 6,54, 54])], ["A", "B"])

from pyspark.sql.functions import *

df.withColumn('Max', expr('aggregate(B, 0L, (a, b) -> if(a < b, b, a))')).show(3, False)

+---+----------------------------------------------+---+
|A  |B                                             |Max|
+---+----------------------------------------------+---+
|1  |[4, 1, 4, 54, 4, 2, 2, 7, 14, 23, 74, 53]     |74 |
|2  |[4, 5, 11, 3, 45, 34, 2, 3, 4]                |45 |
|3  |[4, 0, 32, 23, 23, 5, 23, 2, 37, 8, 6, 54, 54]|54 |
+---+----------------------------------------------+---+

Be aware that the 0L is long type and you should match the type with the element of the array.

this seems to work. Don't know exactly what it does though :-)

def my_max(s):
    return max(s)

from pyspark.sql.types import DateType
my_max2 = F.udf(my_max, DateType())


df.withColumn("mymax", my_max2("B")) 
Related