Отличительная и суммированная агрегация в Spark одной командой

Question

Я немного тренируюсь в Spark и задаюсь вопросом об оптимизации одной из моих задач.Цель проста: рассчитать различное количество заказов и общую стоимость заказа по дате и статусу заказа из следующей таблицы:

Это необходимо сделатьв API Dataframe Spark (Python или Scala) не SQL.

В SQL это было бы просто:

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

Единственный способ заставить его работатьв PySpark выполняется в три этапа:

1. Рассчитать общее количество заказов

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

2. Рассчитать отдельный идентификатор элемента заказа

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

3. Присоединиться к таблицам выше

   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)

Есть ли более разумный путь?Заранее благодарю!

Answer 1

В agg вы можете выполнить оба вычисления в одном groupby, например:

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")
    )