getting maximum value for first element in list of lists in PySpark column

user_7238987 0
Viewed 506

I have a PySpark column with values of list of lists like this

row 1: [['01', '100.0'], ['01', '400.0'], [None, '0.0'], ['06', '0.0'], ['01', '300'], [None, '0.0'], ['06', '200.0']]
row 2: [[None, '200.0'], ['06', '300.0'], ['01', '500'], ['06', '100.0'], ['01', '200'], ['07', '50.0']]

I need to compare elements with same first element in the list of lists and filter out the arrays with the maximum second element for each pair. While the array may have different codes for the first element, I want to filter out array elements containing '01' , '06' or '07' and add two columns to my dataframe.

So the result for a sample row above would look like this:

[['01', '400.0'], ['06', '200.0'], ['07':'0']
[['01', '500.0'], ['06', '300.0'], ['07':'50']

what's the most efficient way to do this?

enter image description here

pyspark apache-spark python
1 Answers

This should do it: 

from pyspark.sql import functions as F

df.withColumn("Max_01",F.when(F.size(F.expr("""filter(arr,x-> exists(x,y->y='01'))"""))!=0,
       F.expr("""array_max(transform(filter(arr, x-> exists(x,y-> y='01')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
  .withColumn("Max_06",F.when(F.size(F.expr("""filter(arr,x-> exists(x,y->y='06'))"""))!=0,
       F.expr("""array_max(transform(filter(arr, x-> exists(x,y-> y='06')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
  .withColumn("Max_07",F.when(F.size(F.expr("""filter(arr,x-> exists(x,y->y='07'))"""))!=0,
       F.expr("""array_max(transform(filter(arr, x-> exists(x,y-> y='07')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
   .show(truncate=False)

#+---------------------------------------------------------------------------------+------+------+------+
#|arr                                                                              |Max_01|Max_06|Max_07|
#+---------------------------------------------------------------------------------+------+------+------+
#|[[01, 100.0], [01, 400.0], [, 0.0], [06, 0.0], [01, 400.0], [, 0.0], [06, 200.0]]|400.0 |200.0 |0.0   |
#|[[, 200.0], [06, 300.0], [01, 500], [06, 100.0], [01, 200], [07, 50.0]]          |500.0 |300.0 |50.0  |
#+---------------------------------------------------------------------------------+------+------+------+
user_101073890
Related