Error in aggregate function Pyspark Dataframe

Viewed 316

I need a help regarding Aggregate function in Pyspark Dataframe. I need to calculate expenses made by customer based on 'buy' or 'sell'.

If buy means I should subtract the amount from the credit limit, if sell means I should add the amount to credit limit

Below is my table

+----------+-----------------+------+----------+----------------+
|account_id|credit_card_limit|amount|      date|transaction_code|
+----------+-----------------+------+----------+----------------+
|     12345|             1000|   400|01/06/2020|             buy|
|     12345|             1000|   100|02/06/2020|             buy|
|     12345|             1000|   500|02/06/2020|            sell|
|     12345|             1000|   200|03/06/2020|             buy|
|     22332|             2000|  1000|02/06/2020|             buy|
|     22332|             2000|   200|03/06/2020|             buy|
+----------+-----------------+------+----------+----------------+

I tried a code but it didn't give me correct results.Below is my code

w = Window.partitionBy(f.lit(0)).orderBy('date')
finaldf=df.groupBy('account_id','credit_card_limit','date').agg(f.sum(f.when(f.col('transaction_code')=='buy',-f.col('amount')).\
              otherwise(f.col('amount'))).alias('expenses')).\
    select('*',(f.col('credit_card_limit')+f.sum(f.col('expenses')).over(w)).alias('credit_left'))

The output I got:

    +----------+-----------------+----------+--------+-----------+
    |account_id|credit_card_limit|      date|expenses|credit_left|
    +----------+-----------------+----------+--------+-----------+
    |     12345|             1000|01/06/2020|    -400|        600|
    |     12345|             1000|02/06/2020|     400|          0|
    |     12345|             1000|03/06/2020|    -200|       -400|
    |     22332|             2000|02/06/2020|   -1000|       1000|
    |     22332|             2000|03/06/2020|    -200|        800|
    +----------+-----------------+----------+--------+-----------+

Here as you can see the credit_left column has wrong answers.

Expected Output:

    +----------+-----------------+----------+--------+-----------+
    |account_id|credit_card_limit|      date|expenses|credit_left|
    +----------+-----------------+----------+--------+-----------+
    |     12345|             1000|01/06/2020|    -400|        600|
    |     12345|             1000|02/06/2020|     400|       1000|
    |     12345|             1000|03/06/2020|    -200|        800|
    |     22332|             2000|02/06/2020|   -1000|       1000|
    |     22332|             2000|03/06/2020|    -200|        800|
    +----------+-----------------+----------+--------+-----------+

I also need to make to credit_left to credit_card_limit incase if the value exceeds the credit_limit.Please help me to solve this problem. Thanks a lot !!

2 Answers

I have assumed that for account 22332 for date 03/06/2020 credicardlimit is 1000 as per logic and expected answer. Please try this and let me know if it works.

df = spark.sql("""
select 12345 as account_id, 1000 as credit_card_limit, 400 as amount, '01/06/2020' as date, 'buy' as  transaction_code
union                                                                                                                                                                                                   
select 12345 as account_id, 1000 as credit_card_limit, 100 as amount, '02/06/2020' as date, 'buy' as  transaction_code
union                                                                                                                                                                                                   
select 12345 as account_id, 1000 as credit_card_limit, 500 as amount, '02/06/2020' as date, 'sell' as  transaction_code
union                                                                                                                                                                                                   
select 12345 as account_id, 1000 as credit_card_limit, 200 as amount, '03/06/2020' as date, 'buy' as  transaction_code
union                                                                                                                                                                                                   
select 22332 as account_id, 2000 as credit_card_limit, 1000 as amount, '02/06/2020' as date, 'buy' as  transaction_code
union
select 22332 as account_id, 1000 as credit_card_limit, 200 as amount, '03/06/2020' as date, 'buy' as  transaction_code
""").orderBy("account_id","date")

df.show()
# source data
# +----------+-----------------+------+----------+----------------+
# |account_id|credit_card_limit|amount|      date|transaction_code|
# +----------+-----------------+------+----------+----------------+
# |     12345|             1000|   400|01/06/2020|             buy|
# |     12345|             1000|   100|02/06/2020|             buy|
# |     12345|             1000|   500|02/06/2020|            sell|
# |     12345|             1000|   200|03/06/2020|             buy|
# |     22332|             2000|  1000|02/06/2020|             buy|
# |     22332|             1000|   200|03/06/2020|             buy|
# +----------+-----------------+------+----------+----------------+

df.createOrReplaceTempView("tmp1")

data1 = spark.sql("""select  account_id,
        credit_card_limit,
        amount, 
        date,
        transaction_code,
        lead(amount) over(partition by account_id order by date) as lead_amt,
        case when transaction_code = 'buy' then -1 * amount else amount end as amount_modified 
from tmp1
order by account_id,date
""")
data1.show()
# +----------+-----------------+------+----------+----------------+--------+---------------+
# |account_id|credit_card_limit|amount|      date|transaction_code|lead_amt|amount_modified|
# +----------+-----------------+------+----------+----------------+--------+---------------+
# |     12345|             1000|   400|01/06/2020|             buy|     100|           -400|
# |     12345|             1000|   100|02/06/2020|             buy|     500|           -100|
# |     12345|             1000|   500|02/06/2020|            sell|     200|            500|
# |     12345|             1000|   200|03/06/2020|             buy|    null|           -200|
# |     22332|             2000|  1000|02/06/2020|             buy|     200|          -1000|
# |     22332|             1000|   200|03/06/2020|             buy|    null|           -200|
# +----------+-----------------+------+----------+----------------+--------+---------------+

data1.createOrReplaceTempView("tmp2")

data2 = spark.sql("""
select account_id,
        credit_card_limit,
        date,
        sum(amount_modified) as expenses,
        case when (credit_card_limit + sum(amount_modified)) > credit_card_limit 
             then credit_card_limit else (credit_card_limit + sum(amount_modified)) 
        end as credit_left
from tmp2
group by account_id, credit_card_limit, date 
order by account_id, date
""")

data2.show()

# +----------+-----------------+----------+--------+-----------+
# |account_id|credit_card_limit|      date|expenses|credit_left|
# +----------+-----------------+----------+--------+-----------+
# |     12345|             1000|01/06/2020|    -400|        600|
# |     12345|             1000|02/06/2020|     400|       1000|
# |     12345|             1000|03/06/2020|    -200|        800|
# |     22332|             2000|02/06/2020|   -1000|       1000|
# |     22332|             1000|03/06/2020|    -200|        800|
# +----------+-----------------+----------+--------+-----------+

I think you need to change the window to :

w = Window.partitionBy(f.col("account_id")).orderBy('date')

then your code works:

w = Window.partitionBy(f.col("account_id")).orderBy('date')

finaldf = (df.groupBy('account_id','credit_card_limit','date')
                .agg(f.sum(f.when(f.col('transaction_code')=='buy',-f.col('amount'))
              .otherwise(f.col('amount'))).alias('expenses')).
    select('*',(f.col('credit_card_limit')+f.sum(f.col('expenses')).over(w))
                                                      .alias('credit_left')))
finaldf.show()

finaldf.show()

+----------+-----------------+----------+--------+-----------+
|account_id|credit_card_limit|      date|expenses|credit_left|
+----------+-----------------+----------+--------+-----------+
|     12345|             1000|01/06/2020|    -400|        600|
|     12345|             1000|02/06/2020|     400|       1000|
|     12345|             1000|03/06/2020|    -200|        800|
|     22332|             2000|02/06/2020|   -1000|       1000|
|     22332|             2000|03/06/2020|    -200|        800|
+----------+-----------------+----------+--------+-----------+
Related