Выберите строку и отобразите имя столбца на основе максимального значения в pyspark

#dataframe #apache-spark #pyspark #apache-spark-sql

#фрейм данных #apache-spark #pyspark #apache-spark-sql

Вопрос:

У меня есть фрейм данных Pyspark как

 | ID|colA|colB|colC|
 --- ---- ---- ---- 
|ID1|   3|   5|   6|
|ID2|   4|  12|   7|
|ID3|   2|  20|  22|
 --- ---- ---- ---- 
 

Я хочу выбрать строку ID3 и выбрать максимальное значение из трех столбцов, а затем отобразить имя столбца с максимальным значением. Поэтому, если я выберу максимальное значение из трех столбцов для строки ID3 , оно должно вернуться, как показано ниже :

 |colC|
 ---- 
|22  |
 ---- 
 

Итак, мой вопрос в том, как нам выбрать одну строку и выбрать имя столбца на основе максимального значения столбца этой выбранной строки из pyspark DataFrame?

Ответ №1:

Вы можете фильтровать ID = 'ID3' и использовать greatest функцию в списке структур, чтобы получить максимальное имя столбца и его значение. Затем поверните, чтобы получить его как столбец:

 from pyspark.sql import functions as F

df1 = df.filter("ID = 'ID3'").withColumn(
    "max_value",
    F.greatest(
        *[F.struct(F.col(c).alias("v"), F.lit(c).alias("name")) for c in df.columns[1:]]
    )
).select("max_value.*").groupBy().pivot("name").agg(F.first("v"))

df1.show(truncate=False)

# ---- 
#|colC|
# ---- 
#|22  |
# ----