python - Pyspark根据对应列中满足特定条件的多列的最大值创建列

Question

假设我有一个 Pyspark 数据框，其中包含 id 和代表代码桶的 3 列。

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

和 3 列表示相应代码桶的数量。

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

这是我正在尝试做的伪代码。

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)))

我将如何做到这一点？

这是一个数据框示例:

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

我正在尝试获取 max_01、max_07 和 max_06 列

Answer 1

对于spark2.4+，可以试试这个。

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     |
#+----------+------+------+------+-----+-----+-----+------+------+------+