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，对于00234， 、和DeviceID中的最大值为 5.6。同样，对于00236， 、和中的最大值为 4.5。 我想根据每个 的最大值检索值。显然对于00234 来说是的。ABCDeviceIDABC
TimeStampDeviceIDDeviceID11-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| 
+--------+----------------+---------+