How to count Total Price in dataframe - dataframe

Question

I have retail data from which I created retail dataframe
spark.sparkContext.addFile('https://raw.githubusercontent.com/databricks/Spark-The-Definitive-Guide/master/data/retail-data/all/online-retail-dataset.csv')
retail_df = spark.read.csv(SparkFiles.get('online-retail-dataset.csv'), header=True, inferSchema=True)\
.withColumn('OverallItems', struct('StockCode', 'Description', 'UnitPrice', 'Quantity', 'InvoiceDate','CustomerID', 'Country'))
then I created retail_array that has two columns InvoiceNo and Items
retail_array = retail_df.groupBy('InvoiceNo')\
.agg(collect_list(col('OverallItems')).alias('Items'))
I want to count total price of invoice items and add to into items column in retail_array.
So far I have written this code:
transformer = lambda x: struct(x['UnitPrice'], x['Quantity'], x['UnitPrice'] * x['Quantity']).cast("struct<UnitPrice:double,Quantity:double,TotalPrice:double>")
TotalPrice_df = retail_array\
.withColumn('TotalPrice', transform("items", transformer))
TotalPrice_df.show(truncate=False)
But with this code Im adding to retail_arraynew column, but I want this new column to be part of items column inretail_array`.
for one invoice item output is like:
--+
|InvoiceNo|Items|TotalPrice |
+---------+---------------------------------------------------------------------------------------
|536366 |[{22633, HAND WARMER UNION JACK, 1.85, 6, 12/1/2010 8:28, 17850, United Kingdom}, {22632, HAND WARMER RED POLKA DOT, 1.85, 6, 12/1/2010 8:28, 17850, United Kingdom}] |[{1.85, 6.0, 11.100000000000001}, {1.85, 6.0, 11.100000000000001}]
I want it count 11.100000000000001 + 11.100000000000001 and add it into items column with no extra column. Also for other invoice items there are sometimes more than two total price I want to add to each other.

Answer 1

Use aggregate instead of transform function to calculate the total price like this:
from pyspark.sql import functions as F
retail_array = retail_df.groupBy("InvoiceNo").agg(
F.collect_list(F.col("OverallItems")).alias("Items")
).withColumn(
"TotalPrice",
F.aggregate("items", F.lit(.0), lambda acc, x: acc + (x["Quantity"] * x["UnitPrice"]))
)
Note however that you can actually calculate this TotalPrice in the same aggregation when you collect the list of structs and thus avoid, additional calculations by iterating on array elements:
retail_array = retail_df.groupBy("InvoiceNo").agg(
F.collect_list(F.col("OverallItems")).alias("Items"),
F.sum(F.col("Quantity") * F.col("UnitPrice")).alias("TotalPrice")
)
retail_array.show(1)
#+---------+--------------------+------------------+
#|InvoiceNo| Items| TotalPrice|
#+---------+--------------------+------------------+
#| 536366|[{22633, HAND WAR...|22.200000000000003|
#+---------+--------------------+------------------+
But with this code I'm adding to retail_array new column, but I want this new column to be part of items column in retail_array
Note sure I correctly understood this part. Items column is an array of structs, that does not make much sens to replicate the total price of an InvoiceNo in each
of its items.
That said, if you really want to do this, you can use transform after calculating the total price (step above):
result = retail_array.withColumn(
"Items",
F.transform("Items", lambda x: x.withField("TotalPrice", F.col("TotalPrice")))
).drop("TotalPrice")
result.show(1, False)
#+---------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
#|InvoiceNo|Items |
#+---------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
#|536366 |[{22633, HAND WARMER UNION JACK, 1.85, 6, 12/1/2010 8:28, 17850, United Kingdom, 22.200000000000003}, {22632, HAND WARMER RED POLKA DOT, 1.85, 6, 12/1/2010 8:28, 17850, United Kingdom, 22.200000000000003}]|
#+---------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+