Вычисление разницы значений между двумя датами с помощью pyspark

#sql #pyspark

Вопрос:

Мои данные таковы:

год_месяц идентификатор пользователя просмотры страниц посещения
2020-03 2 8 3
2021-03 27 4 3
2021-05 23 75 7
2020-05 23 17 7
2020-08 339 253 169
2020-08 892 31 4
2021-08 339 4 3

И я хотел сгруппировать по месяцам года, рассчитав разницу просмотров страниц и посещений с одного года(2020) до следующего(2021).

Итак, я думал, что вывод должен быть чем-то похожим (без содержимого в круглых скобках).:

последний месяц разница(просмотры страниц) разница(посещений)
2021-03 -4(4-8) 0(3-3)
2021-05 58(75-17) 0(7-7)
2021-08 -280(4-284) -170(3-173)

Но я не уверен, как это сделать с векторизацией, я думал передать это пандам и сделать это с помощью цикла for, но хотел узнать, как выполнять такого рода действия векторизированным способом с помощью pyspark или sparksql, что, я думаю, они будут намного быстрее.

Ответ №1:

Основная идея заключается в использовании функции окна для сравнения месяцев. Проверьте мои комментарии для получения дополнительных объяснений

 from pyspark.sql import functions as F
from pyspark.sql import Window as W

(df
    # since you'd want to compare month and year separately,
    # we have to separate them out using split function
    .withColumn('year', F.split('year_month', '-')[0].cast('int'))
    .withColumn('month', F.split('year_month', '-')[1].cast('int'))
 
    # you have multiple rows per year_month
    # so we have to group and sum the similar records
    .groupBy('year', 'month')
    .agg(
        F.sum('pageviews').alias('pageviews'),
        F.sum('visits').alias('visits')
    )
    
    # now, you need to compare 2021's months with 2020's months,
    # you'd have to use lag window function, pay attention to the orderBy window
    .withColumn('prev_pageviews', F.lag('pageviews').over(W.orderBy('month', 'year')))
    .withColumn('prev_visits', F.lag('visits').over(W.orderBy('month', 'year')))

    # with current pageviews/visits and previous pageviews/visits on the same row
    # you can easily calculate the difference between months
    .withColumn('diff_pageviews', F.col('pageviews') - F.col('prev_pageviews'))
    .withColumn('diff_visits', F.col('visits') - F.col('prev_visits'))

    # select only necessary colums and rows
    .select('year', 'month', 'diff_pageviews', 'diff_visits')
    .where(F.col('year') == 2021)

    .show()
)

# Output
#  ---- ----- -------------- ----------- 
# |year|month|diff_pageviews|diff_visits|
#  ---- ----- -------------- ----------- 
# |2021|    3|            -4|          0|
# |2021|    5|            58|          0|
# |2021|    8|          -280|       -170|
#  ---- ----- -------------- -----------