Pyspark create column based on maximum of multiple columns that match a certain condition in corresponding columns

user_7238987 0
Viewed 576

Let's say I have a Pyspark dataframe with id and 3 columns representing code buckets.

col_buckets ["code_1", "code_2", "code_3"]

and 3 columns representing amounts for corresponding code buckets.

amt_buckets = ["code_1_amt", "code_2_amt", "code_3_amt" ]

Here is a pseudocode for what I am trying to do.

for el in ['01', '06', '07']
    df= df.withColumn("max_amt_{el}", max(df.select(max(**amt_buckets**) for corresponding col_indices of amt_buckets if ***any of col_buckets*** ==el)))

how would I accomplish this?

here is a dataframe example for this:

enter image description here

Primary_id  Code_1  Code_2  Code_3  Amt_1   Amt_2   Amt_3   Max_01  Max_07  Max_06
Xxxxx998    Null    01      04      2000    1000    100     1000    0       0
Xxxxx997    01      01      07      200     300     400     300     400     0
Xxxxx996    07      Null    Null    100     Null    Null    0       100     0
Xxxx910     Null    Null    Null    300     100     200     0       0       0

I am trying to get the max_01, max_07 and max_06 columns

pyspark apache-spark python
1 Answers

For spark2.4+, you can try this.

df.show() #sample dataframe
#+----------+------+------+------+-----+-----+-----+
#|Primary_id|Code_1|Code_2|Code_3|Amt_1|Amt_2|Amt_3|
#+----------+------+------+------+-----+-----+-----+
#|  Xxxxx998|  null|    01|    04| 2000| 1000|  100|
#|  Xxxxx997|    01|    01|    07|  200|  300|  400|
#|  Xxxxx996|    07|  null|  null|  100| null| null|
#|   Xxxx910|  null|  null|  null|  300|  100|  200|
#+----------+------+------+------+-----+-----+-----+

from pyspark.sql import functions as F

dictionary = dict(zip(['Code_1','Code_2','Code_3'], ['Amt_1','Amt_2','Amt_3']))

df.withColumn("trial", F.array(*[F.array(F.col(x),F.col(y).cast("string"))\
                                          for x,y in dictionary.items()]))\
  .withColumn("Max_01",F.when(F.size(F.expr("""filter(trial,x-> exists(x,y->y='01'))"""))!=0,\
       F.expr("""array_max(transform(filter(trial, x-> exists(x,y-> y='01')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
  .withColumn("Max_06",F.when(F.size(F.expr("""filter(trial,x-> exists(x,y->y='06'))"""))!=0,\
       F.expr("""array_max(transform(filter(trial, x-> exists(x,y-> y='06')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
  .withColumn("Max_07",F.when(F.size(F.expr("""filter(trial,x-> exists(x,y->y='07'))"""))!=0,\
       F.expr("""array_max(transform(filter(trial, x-> exists(x,y-> y='07')),z-> float(z[1])))"""))\
             .otherwise(F.lit(0)))\
  .drop("trial").show(truncate=False)

#+----------+------+------+------+-----+-----+-----+------+------+------+
#|Primary_id|Code_1|Code_2|Code_3|Amt_1|Amt_2|Amt_3|Max_01|Max_07|Max_06|
#+----------+------+------+------+-----+-----+-----+------+------+------+
#|Xxxxx998  |null  |01    |04    |2000 |1000 |100  |1000  |0     |0     |
#|Xxxxx997  |01    |01    |07    |200  |300  |400  |300   |400   |0     |
#|Xxxxx996  |07    |null  |null  |100  |null |null |0     |100   |0     |
#|Xxxx910   |null  |null  |null  |300  |100  |200  |0     |0     |0     |
#+----------+------+------+------+-----+-----+-----+------+------+------+
user_101073890
Related