pyspark conditional cumulative sum

user_19701388 0
Viewed 71

I have a pyspark dataframe with two dates - bill and payment date. I want to create a column that has the sum of amounts of bills billed and paid before the bill date of that row. Also, this needs to be done for every buyer ID individually. Example:

Buyer Bill date Payment Date Amount New Column
1 2021-02-02 2021-02-20 100 0
1 2021-03-02 2021-03-10 400 100
1 2021-04-02 2021-05-25 500 500
1 2021-05-02 2021-06-03 300 500
1 2021-06-02 2021-07-20 200 1000
2 2021-04-10 2021-05-25 1000 0
2 2021-05-11 2021-06-03 3000 0
2 2021-06-15 2021-07-20 2000 4000

Pandas equivalent of what I am looking for would be:

def to_value(row):
    return dt[(dt['pay_dt']<row['pay_dt'])&(dt['pay_dt']<row['bill_dt'])&(dt['buyer_id']==row['buyer_id'])].amount.sum()

dt['new_col']=dt.apply(to_value,axis=1)
pyspark pandas python
1 Answers

You can use pandas_udf() and do conditional processing there:

import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql import SparkSession, Window
from pyspark.sql.types import IntegerType

def conditional_sum(data: pd.DataFrame) -> int:
    df = data.apply(pd.Series)  # transform dict into separate columns
    return df.loc[df['Bill date'].max() > df['Payment Date']]['Amount'].sum()

udf_conditional_sum = F.pandas_udf(conditional_sum, IntegerType())

w = Window.partitionBy("Buyer").orderBy("Bill date").rowsBetween(Window.unboundedPreceding, Window.currentRow)

(
    df
    .withColumn("Conditional sum", udf_conditional_sum(F.struct("Bill date", "Payment Date", "Amount")).over(w))
    .show(truncate=False)
)
user_46489690
Related