python - PySpark:根据其他三列的最大值查找一列的值

我是新手。我使用的数据帧如下

DeviceID       TimeStamp           A      B     C
00234       11-03-2014 05:55      5.6    2.3   3.3
00235       11-03-2014 05:33      2.8    0.9   4.2
00236       11-03-2014 06:15      3.5    0.1   1.3
00234       11-03-2014 07:23      2.5    0.2   3.9
00236       11-03-2014 07:33      2.5    4.5   2.9

从上述样本可以看出，对于df00234，df、DeviceID和A之间的最大值为5.6。同样，对于B00236，C、DeviceID和A之间的最大值为4.5。
我想根据每个B的最大值检索C值。显然，对于TimeStamp00234，它是DeviceID。
虽然我没有尝试过任何方法，但是，下面的方法是否有效？

from pyspark.sql import function as F
max_value = df.groupby('DeviceID').agg(F.greatest('A','B','C').alias('max_value'))
df.withColumn('Max-TimeStamp',where(# please help me in putting the right codes))

结果DeviceID应如下所示

DeviceID    Max_Value     Max-TimeStamp
00234          5.6        11-03-2014 05:55
00236          4.5        11-03-2014 07:33

任何帮助都将不胜感激。谢谢。

最佳答案

您可以通过Window函数来实现这一点：

import pyspark.sql.functions as F
from pyspark.sql import Window

l = [('00234'      , '11-03-2014 05:55',      5.6 ,   2.3 ,  3.3),
     ('00235'      , '11-03-2014 05:33'   ,   2.8,    0.9  , 4.2),
     ('00236'      , '11-03-2014 06:15'  ,    3.5 ,   0.1  , 1.3),
     ('00234'      , '11-03-2014 07:23' ,     2.5  ,  0.2 ,  3.9),
     ('00236'      , '11-03-2014 07:33',      2.5   , 4.5,   2.9)]

columns = ['DeviceID', 'TimeStamp', 'A','B','C']

df=spark.createDataFrame(l, columns)

w = Window.partitionBy('DeviceID')

df = df.select('DeviceID', 'TimeStamp', F.greatest('A','B','C').alias('max_value'))

df.withColumn('bla', F.max('max_value').over(w)).where(F.col('max_value') == F.col('bla')).drop('bla').show()

输出：

+--------+----------------+---------+
|DeviceID| TimeStamp      |max_value|
+--------+----------------+---------+
|   00236|11-03-2014 07:33|      4.5|
|   00234|11-03-2014 05:55|      5.6|
|   00235|11-03-2014 05:33|      4.2|
+--------+----------------+---------+