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

Question

我正在使用数据框 df 如下

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

从上面的示例df可以看出，对于DeviceID 00234，A、B中的最大值> 和 C 是 5.6。同样，对于 DeviceID 00236，A、B 和 C 中的最大值为 4.5。
我想根据每个 DeviceID 的最大值检索 TimeStamp 值。显然，对于 DeviceID 00234，它是 11-03-2014 05:55。

虽然我没有尝试过任何方法，但是，以下方法是否有效？

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

生成的 df 应该如下所示

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

Answer 1

您可以使用 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| 
+--------+----------------+---------+