Início / coding / Perguntas / 79084670
Accepted
John
Asked: 2024-10-14 12:01:55 +0800 CST

Criar coluna de saldo atual no PySpark

  • 772

Tenho o dataframe abaixo criado no código PySpark:

+---------------+-------------+---------------+------+
|TransactionDate|AccountNumber|TransactionType|Amount|
+---------------+-------------+---------------+------+
|     2023-01-01|          100|         Credit|  1000|
|     2023-01-02|          100|         Credit|  1500|
|     2023-01-03|          100|          Debit|  1000|
|     2023-01-02|          200|         Credit|  3500|
|     2023-01-03|          200|          Debit|  2000|
|     2023-01-04|          200|         Credit|  3500|
|     2023-01-13|          300|         Credit|  4000|
|     2023-01-14|          300|          Debit|  4500|
|     2023-01-15|          300|         Credit|  5000|
+---------------+-------------+---------------+------+

Preciso imprimir outra coluna como CurrentBalance .

Saída esperada:

+---------------+-------------+---------------+------+--------------+
|TransactionDate|AccountNumber|TransactionType|Amount|CurrentBalance|
+---------------+-------------+---------------+------+--------------+
|     2023-01-01|          100|         Credit|  1000|          1000|
|     2023-01-02|          100|         Credit|  1500|          2500|
|     2023-01-03|          100|          Debit|  1000|          1500|
|     2023-01-02|          200|         Credit|  3500|          3500|
|     2023-01-03|          200|          Debit|  2000|          1500|
|     2023-01-04|          200|         Credit|  3500|          5000|
|     2023-01-13|          300|         Credit|  4000|          4000|
|     2023-01-14|          300|          Debit|  4500|          -500|
|     2023-01-15|          300|         Credit|  5000|          1000|
+---------------+-------------+---------------+------+--------------+

Tentei usar a data mínima e passar a data na condição when para calcular o crédito e o débito, mas parece que não está funcionando.

# Find minimum date in TransactionDate column, grouped by AccountNumber column
df_new.groupBy('AccountNumber').agg(f.min('TransactionDate').alias('min_date'))
dataframe

1 respostas

  1. Best Answer

    Você precisará de uma função de janela para isso. Funções de janela fazem cálculos para cada linha em uma partição (grupo). Nesse caso, você precisa somar todos os valores linha após linha.

    Além disso, a soma simples não funcionaria, pois você não tem números negativos, então você deve fazê-los usando a entrada da coluna "TransactionType".

    Dados de exemplo:

    from pyspark.sql import functions as F, Window as W

df = spark.createDataFrame(
    [('2023-01-01', '100', 'Credit', 1000),
     ('2023-01-02', '100', 'Credit', 1500),
     ('2023-01-03', '100', 'Debit', 1000),
     ('2023-01-02', '200', 'Credit', 3500),
     ('2023-01-03', '200', 'Debit', 2000),
     ('2023-01-04', '200', 'Credit', 3500),
     ('2023-01-13', '300', 'Credit', 4000),
     ('2023-01-14', '300', 'Debit', 4500),
     ('2023-01-15', '300', 'Credit', 5000)],
    ['TransactionDate', 'AccountNumber', 'TransactionType', 'Amount'])

    Roteiro:

    sign = F.when(F.col('TransactionType') == 'Debit', -1).otherwise(1)
amount = sign * F.col('amount')
window = W.partitionBy('AccountNumber').orderBy('TransactionDate')

df = df.withColumn('CurrentBalance', F.sum(amount).over(window))

df.show() 
# +---------------+-------------+---------------+------+--------------+
# |TransactionDate|AccountNumber|TransactionType|Amount|CurrentBalance|
# +---------------+-------------+---------------+------+--------------+
# |     2023-01-01|          100|         Credit|  1000|          1000|
# |     2023-01-02|          100|         Credit|  1500|          2500|
# |     2023-01-03|          100|          Debit|  1000|          1500|
# |     2023-01-02|          200|         Credit|  3500|          3500|
# |     2023-01-03|          200|          Debit|  2000|          1500|
# |     2023-01-04|          200|         Credit|  3500|          5000|
# |     2023-01-13|          300|         Credit|  4000|          4000|
# |     2023-01-14|          300|          Debit|  4500|          -500|
# |     2023-01-15|          300|         Credit|  5000|          4500|
# +---------------+-------------+---------------+------+--------------+
    • 0

relate perguntas