Creating Total and percentage of total columns in Pyspark

user_10523971 0
Viewed 170

Here is my test data

test = spark.createDataFrame([
    ("2018-06-03",2, 4, 4 ),
    ("2018-06-04",4, 3, 3 ),
    ( "2018-06-03",8, 1, 1),
    ("2018-06-01",3, 1, 1),
    ( "2018-06-05", 3, 2, 0),
])\
  .toDF( "transactiondate", "SalesA", "SalesB","SalesC")
test.show()

I would like to add a row-wise total column and % of the total column corresponding to each sales category (A, B and C)

Desired Output:

+---------------+------+------+------+----------+------+------+------+
|transactiondate|SalesA|SalesB|SalesC|TotalSales|Perc_A|Perc_B|Perc_C|
+---------------+------+------+------+----------+------+------+------+
|     2018-06-03|     2|     4|     4|        10|   0.2|   0.4|   0.4|
|     2018-06-04|     4|     3|     3|        10|   0.4|   0.3|   0.3|
|     2018-06-03|     8|     1|     1|        10|   0.8|   0.1|   0.1|
|     2018-06-01|     3|     1|     1|         5|   0.6|   0.2|   0.2|
|     2018-06-05|     3|     2|     0|         5|   0.6|   0.4|   0.0|
+---------------+------+------+------+----------+------+------+------+

How can I do it in pyspark?

Edit: I want the code to be adaptable even if I add more items, i.e. if I have one more column salesD, code should create total and percentage columns. (i.e. columns shouldn't be hardcoded)

pyspark apache-spark-sql apache-spark
3 Answers

You can use selectExpr and do simple arithmetic SQL operations for each added columns

test = test.selectExpr("*", 
                       "SalesA+SalesB+SalesC as TotalSales",
                       "SalesA/(SalesA+SalesB+SalesC) as Perc_A",
                       "SalesB/(SalesA+SalesB+SalesC) as Perc_B",
                       "SalesC/(SalesA+SalesB+SalesC) as Perc_C"
                       )

or use a more flexible solution

from pyspark.sql.functions import col, expr

# columns to be included in TotalSales calculation
cols = ['SalesA', 'SalesB', 'SalesC']

test = (test
        .withColumn('TotalSales', expr('+'.join(cols)))
        .select(col('*'), 
                *[expr('{0}/TotalSales {1}'.format(c,'Perc_'+c)) for c in cols]))
user_94778430

One option is to use several withColumn statements

import pyspark.sql.functions as F

test\
  .withColumn('TotalSales', F.col('SalesA') + F.col('SalesB') + F.col('SalesC'))\
  .withColumn('Perc_A', F.col('SalesA') / F.col('TotalSales'))\
  .withColumn('Perc_B', F.col('SalesB') / F.col('TotalSales'))\
  .withColumn('Perc_C', F.col('SalesC') / F.col('TotalSales'))
user_74654620

Try this spark-sql solution

test.createOrReplaceTempView("sales_table")

sales=[ x for x in test.columns if x.upper().startswith("SALES") ]
sales2="+".join(sales)
print(str(sales)) # ['SalesA', 'SalesB', 'SalesC']

per_sales=[ x +"/TotalSales as " + "Perc_" +x  for x in sales ]
per_sales2=",".join(per_sales)
print(str(per_sales)) # ['SalesA/TotalSales as Perc_SalesA', 'SalesB/TotalSales as Perc_SalesB', 'SalesC/TotalSales as Perc_SalesC']

spark.sql(f"""
with t1 ( select *, {sales2} TotalSales from sales_table )
select *, {per_sales2} from t1 
""").show(truncate=False)

+---------------+------+------+------+----------+-----------+-----------+-----------+
|transactiondate|SalesA|SalesB|SalesC|TotalSales|Perc_SalesA|Perc_SalesB|Perc_SalesC|
+---------------+------+------+------+----------+-----------+-----------+-----------+
|2018-06-03     |2     |4     |4     |10        |0.2        |0.4        |0.4        |
|2018-06-04     |4     |3     |3     |10        |0.4        |0.3        |0.3        |
|2018-06-03     |8     |1     |1     |10        |0.8        |0.1        |0.1        |
|2018-06-01     |3     |1     |1     |5         |0.6        |0.2        |0.2        |
|2018-06-05     |3     |2     |0     |5         |0.6        |0.4        |0.0        |
+---------------+------+------+------+----------+-----------+-----------+-----------+

You can also use the aggregate() higher order function to sum the sales* columns. But for this the columns must be of Integer/double type, not long.

test2=test.withColumn("SalesA",expr("cast(salesa as int)"))\
      .withColumn("SalesB",expr("cast(salesb as int)"))\
      .withColumn("SalesC",expr("cast(salesc as int)"))
test2.createOrReplaceTempView("sales_table2")

sales3=",".join(sales)  # just join the sales columns with comma

spark.sql(f"""
with t1 ( select *,  aggregate(array({sales3}),0,(acc,x) -> acc+x) TotalSales from sales_table2 )
select *, {per_sales2} from t1 
""").show(truncate=False)

+---------------+------+------+------+----------+-----------+-----------+-----------+
|transactiondate|SalesA|SalesB|SalesC|TotalSales|Perc_SalesA|Perc_SalesB|Perc_SalesC|
+---------------+------+------+------+----------+-----------+-----------+-----------+
|2018-06-03     |2     |4     |4     |10        |0.2        |0.4        |0.4        |
|2018-06-04     |4     |3     |3     |10        |0.4        |0.3        |0.3        |
|2018-06-03     |8     |1     |1     |10        |0.8        |0.1        |0.1        |
|2018-06-01     |3     |1     |1     |5         |0.6        |0.2        |0.2        |
|2018-06-05     |3     |2     |0     |5         |0.6        |0.4        |0.0        |
+---------------+------+------+------+----------+-----------+-----------+-----------+
user_68670480
Related