Pyspark: как создать окно над датами

Question

Я хочу выполнить расчет только для определенного подмножества информационного кадра, создав окно, которое может включать данную дату:

df=df.orderBy('Date')
window_row = Window.partitionBy('I1','Id2')
df=df.withColumn('max_Date', when((col('Date')<=target_date),max('Date').over(window_row)))
df=df.withColumn('cum_sum', when((col('Date')==col('max_Date')),sum('Sale').over(window_row)))

Когда target_date= '2020-01-01', я получаю следующий вывод:

|I1| Id2|         Date|  Sale|  max_Date|cum_sum|
+--+----+-------------+------+----------+------+
|AA|  B0|   2019-07-01|     1|2020-12-01|   null|
|AA|  B0|   2020-01-01|    23|2020-12-01|   null|
|AA|  B0|   2020-01-01|     2|2020-01-01|   null|
|AA|  B0|   2020-02-01|     0|      null|   null|
|AA|  B0|   2020-12-01|   116|      null|   null|
|BB|  C0|   2019-03-01|     1|2020-03-01|   null|
|BB|  C0|   2019-05-01|    26|2020-03-01|   null|
|BB|  C0|   2020-03-01|     1|      null|   null|
|CC|  B0|   2019-03-01|     8|2019-04-01|   null|
|CC|  B0|   2019-04-01|     1|2019-04-01|      1|

Тем не менее, желаемый результат:

|I1| Id2|         Date|  Sale|  max_Date|cum_sum|
+--+----+-------------+------+----------+------+
|AA|  B0|   2019-07-01|     1|2020-01-01|   null|
|AA|  B0|   2020-01-01|    23|2020-01-01|     25|
|AA|  B0|   2020-01-01|     2|2020-01-01|     25|
|AA|  B0|   2020-02-01|     0|      null|   null|
|AA|  B0|   2020-12-01|   116|      null|   null|
|BB|  C0|   2019-03-01|     1|2019-05-01|   null|
|BB|  C0|   2019-05-01|    26|2019-05-01|     26|
|BB|  C0|   2020-03-01|     1|      null|   null|
|CC|  B0|   2019-03-01|     8|2019-04-01|   null|
|CC|  B0|   2019-04-01|     1|2019-04-01|      1|

Как мне эффективно это реализовать?

Answer 1

Это может быть полезно

>>> w = Window.partitionBy("date")
>>> w_preceding = Window.partitionBy("I1","ID2").orderBy('DATE').rangeBetween(Window.unboundedPreceding, 0)
>>> df=df.withColumn('max_Date', when((col('Date')<=target_date),max('Date').over(w)))
>>> df=df.withColumn('cum_sum', when((col('Date')==col('max_Date')),sum('Sale').over(w_preceding)))
>>> df.show()
+----------+---+---+----+----------+-------+
|      Date| I1|Id2|Sale|  max_Date|cum_sum|
+----------+---+---+----+----------+-------+
|2019-05-01| BB| C0|  26|2019-05-01|   26.0|
|2020-03-01| BB| C0|   1|2020-03-01|   27.0|
|2019-07-01| AA| B0|   1|2019-07-01|    2.0|
|2019-07-01| AA| B0|   1|2019-07-01|    2.0|
|2020-02-01| AA| B0|   0|2020-02-01|    2.0|
|2020-09-01| AA| B0|   0|      null|   null|
|2019-04-01| CC| B0|   1|2019-04-01|    1.0|
+----------+---+---+----+----------+-------+

Answer 2

Вот небольшой пример, который поможет вам выбрать правильное направление для работы с окнами. Вот пример окна с операциями суммы, максимума и накопленной суммы. У меня были некоторые проблемы с отслеживанием данных и желаемым выводом, поэтому я создал небольшую выборку данных. Надеюсь, это поможет.

# create data with pandas
import pandas as pd
data = {
    "DATE": ["2020", "2019", "2018", "2020", "2019", "2018"],
    "VALUE": [1, 3, 4, 5, 9, 6]
}
pd_df  = pd.DataFrame(data)

# make it a spark dataframe
spark_df = spark.createDataFrame(pd_df)
spark_df.show()

'''
Here is the shown input data frame
+----+-----+
|DATE|VALUE|
+----+-----+
|2020|    1|
|2019|    3|
|2018|    4|
|2020|    5|
|2019|    9|
|2018|    6|
+----+-----+
'''

# perform window operations
from pyspark.sql.window import Window
from pyspark.sql.functions import max, sum

w = Window.partitionBy("DATE")
w_preceding = Window.partitionBy("DATE").orderBy('VALUE').rangeBetween(Window.unboundedPreceding, 0)

spark_df = spark_df.withColumn("MAX_WINDOW", max("VALUE").over(w))
spark_df = spark_df.withColumn("SUM_WINDOW", sum("VALUE").over(w))
spark_df = spark_df.withColumn("CUM_SUM_WINDOW", sum("VALUE").over(w_preceding))
spark_df.show()

'''
here is the final result dataframe
+----+-----+----------+----------+--------------+
|DATE|VALUE|MAX_WINDOW|SUM_WINDOW|CUM_SUM_WINDOW|
+----+-----+----------+----------+--------------+
|2020|    1|         5|         6|             1|
|2020|    5|         5|         6|             6|
|2019|    3|         9|        12|             3|
|2019|    9|         9|        12|            12|
|2018|    4|         6|        10|             4|
|2018|    6|         6|        10|            10|
+----+-----+----------+----------+--------------+
'''

Дайте мне знать, если это поможет.