получить максимум из столбца и сравнить с каждым элементом столбца

Question

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

id | value | date1       | date2
-------------------------------------
1  | 20    | 2015-09-01  | 2018-03-01 
1  | 30    | 2019-04-04  | 2015-03-02 
1  | 40    | 2014-01-01  | 2016-06-09 
2  | 15    | 2014-01-01  | 2013-06-01 
2  | 25    | 2019-07-18  | 2016-07-07 

, и я хочу вернуть для каждого id sum(value), где date1<max(date2) для этого id.В приведенном выше примере мы получим:

id | sum_value 
-----------
1  | 60     
2  | 15 

, поскольку для id 1 max(date2) равно 2018-03-01, а первая и третья строки соответствуют условию date1<max(date2), и поэтому значение является суммой20 и 40.

Я попробовал приведенный ниже код, но мы не можем использовать max вне функции agg.

df.withColumn('sum_value',F.when(F.col('date1')<F.max(F.col('date2')), value).otherwise(0))
            .groupby(['id']) 

У вас есть предложения?Таблица состоит из 2 миллиардов строк, поэтому я ищу другие варианты, кроме повторного соединения.

Answer 1

Вы можете использовать функцию Window.Прямой перевод ваших требований будет:

from pyspark.sql.functions import col, max as _max, sum as _sum
from pyspark.sql import Window

df.withColumn("max_date2", _max("date2").over(Window.partitionBy("id")))\
    .where(col("date1") < col("max_date2"))\
    .groupBy("id")\
    .agg(_sum("value").alias("sum_value"))\
    .show()
#+---+---------+
#| id|sum_value|
#+---+---------+
#|  1|     60.0|
#|  2|     15.0|
#+---+---------+