Dataframe 选择一行并根据pyspark中的最大值显示列名

我有一个Pyspark数据帧作为

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

我想选择行

ID3

，选择三列的最大值，然后显示最大值的列名。因此，如果我为row

ID3

选择三列的最大值，它应该返回如下：

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

---

因此，我的问题是，我们如何选择一行并根据pyspark DataFrame中所选行的最大列值选择列名？

您可以在

ID='ID3'

上进行筛选，并在结构列表上使用

grest

函数来获取最大列名及其值。然后通过透视将其作为列：

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