Distinct and sum aggregation in Spark using one command

user_745504 0
Viewed 4377

I am doing some Spark training and are wondering about optimizing one of my tasks. The goal is simple: calculate distinct number of orders and total order value by order date and status from the following table:

Input table

This has to be done in Spark's Dataframe API (Python or Scala), not SQL.

In SQL, it would be simple:

select order_status, order_date, count(distinct order_item_id), sum(order_item_subtotal) 
from df 
group by order_status, order_date

The only way I could make it work in PySpark is in three steps:

  1. Calculate total orders

    df_to = df.groupby('order_date','order_status') \ .agg(countDistinct(df.order_id) \ .alias('total_orders'))

  2. Calculate distinct order item id

    df_ta = df.groupby('order_date','order_status') \ .sum('order_item_subtotal') \ .alias('total_amount') \ .withColumnRenamed("sum(order_item_subtotal)","total_amount")

  3. Join tables above

    dfout = df_to.join(df_ta, [df_to.order_date == df_ta.order_date, df_to.order_status == df_ta.order_status], 'inner').select(df_to.order_date, df_to.order_status, df_to.total_orders, df_ta.total_amount)

Is there a smarter way? Thank in you in advance!

pyspark apache-spark-sql apache-spark scala python
1 Answers

Within agg you can perform both calculations in one groupby like this:

import pyspark.sql.functions as func

df_agg = df.groupby("order_date", "order_status").\
    agg(
        func.countDistinct("order_id").alias("total_orders"),
        func.sum("order_item_subtotal").alias("total_amount")
    )
user_27038020
Related