如何在 Pyspark 中按字段对每个组的相同值求和

Question

I have a dataframe called "df" which is in groupby by the "accountname" field, each entry in this column has a cost that can be the same or different, I need to add it only when it is different. This is the original df:

accountname |   namespace   |   cost    
account001  |   ns1         |   11      
account001  |   ns1         |   11      
account001  |   ns1         |   11      
account001  |   ns1         |   11      
account001  |   ns2         |   10      
account001  |   ns2         |   10      
account002  |   ns3         |   50      
account002  |   ns3         |   50      
account002  |   ns3         |   50      
account003  |   ns4         |   5    

The only entry that has different costs within the "accountname" field is "account001", I only need to add 11 + 10 once. And I need to get something like this:

accountname |   namespace   |   cost    |   cost_to_pay
account001  |   ns1         |   11      |   21
account001  |   ns1         |   11      |   21
account001  |   ns1         |   11      |   21
account001  |   ns1         |   11      |   21
account001  |   ns2         |   10      |   21
account001  |   ns2         |   10      |   21
account002  |   ns3         |   50      |   50
account002  |   ns3         |   50      |   50
account002  |   ns3         |   50      |   50
account003  |   ns4         |   5       |   5

Any idea how to do it? Thanks in advance.

Answer 1

You can use collect_set over window partitioned by accountname to get distinct cost values, then sum the elements of the resulting array using aggregate function:

from pyspark.sql import functions as F

df1 = df.withColumn(
    "cost_to_pay",
    F.expr("aggregate(collect_set(cost) over(partition by accountname), 0D, (acc, x) -> acc + x)")
)

df1.show()
#+-----------+---------+----+-----------+
#|accountname|namespace|cost|cost_to_pay|
#+-----------+---------+----+-----------+
#| account003|      ns4|   5|          5|
#| account001|      ns1|  11|         21|
#| account001|      ns1|  11|         21|
#| account001|      ns1|  11|         21|
#| account001|      ns1|  11|         21|
#| account001|      ns2|  10|         21|
#| account001|      ns2|  10|         21|
#| account002|      ns3|  50|         50|
#| account002|      ns3|  50|         50|
#| account002|      ns3|  50|         50|
#+-----------+---------+----+-----------+

Answer 2

You can remove duplicates using distinct , group by accountname and sum the cost, and join back to the original dataframe using accountname :

import pyspark.sql.functions as F

df2 = (df.dropDuplicates(['accountname', 'namespace', 'cost'])
         .groupBy('accountname')
         .agg(F.sum('cost').alias('cost_to_pay'))
         .join(df, 'accountname')
         .select('accountname', 'namespace', 'cost', 'cost_to_pay')
      )

df2.show()
+-----------+---------+----+-----------+
|accountname|namespace|cost|cost_to_pay|
+-----------+---------+----+-----------+
| account001|      ns1|  11|         21|
| account001|      ns1|  11|         21|
| account001|      ns1|  11|         21|
| account001|      ns1|  11|         21|
| account001|      ns2|  10|         21|
| account001|      ns2|  10|         21|
| account002|      ns3|  50|         50|
| account002|      ns3|  50|         50|
| account002|      ns3|  50|         50|
| account003|      ns4|   5|          5|
+-----------+---------+----+-----------+