sql - 在 Pyspark 数据框上查找迄今为止的月份和月份

Question

我在 Spark 中有以下数据框(使用 PySpark):DT_BORD_REF :时间戳列，COUNTRY_ALPHA :国家 Alpha-3 代码，working_day_flag : 日期是否是该国家/地区的工作日
我需要添加两个字段:

该国家/地区从月初开始的工作日计数(本月至今)

该国家(即将过去的月份)到该月月底剩余的工作日数

好像是一个窗口函数的应用，但我想不通

+-------------------+-------------+----------------+
|        DT_BORD_REF|COUNTRY_ALPHA|working_day_flag|
+-------------------+-------------+----------------+
|2021-01-01 00:00:00|          FRA|               N|
|2021-01-01 00:00:00|          ITA|               N|
|2021-01-01 00:00:00|          BRA|               N|
|2021-01-02 00:00:00|          BRA|               N|
|2021-01-02 00:00:00|          FRA|               N|
|2021-01-02 00:00:00|          ITA|               N|
|2021-01-03 00:00:00|          ITA|               N|
|2021-01-03 00:00:00|          BRA|               N|
|2021-01-03 00:00:00|          FRA|               N|
|2021-01-04 00:00:00|          BRA|               Y|
|2021-01-04 00:00:00|          FRA|               Y|
|2021-01-04 00:00:00|          ITA|               Y|
|2021-01-05 00:00:00|          FRA|               Y|
|2021-01-05 00:00:00|          BRA|               Y|
|2021-01-05 00:00:00|          ITA|               Y|
|2021-01-06 00:00:00|          ITA|               N|
|2021-01-06 00:00:00|          FRA|               Y|
|2021-01-06 00:00:00|          BRA|               Y|
|2021-01-07 00:00:00|          ITA|               Y|
+-------------------+-------------+----------------+

Answer 1

您可以使用 count_if 进行条件计数:

df.createOrReplaceTempView('df')

result = spark.sql("""
select *,
    count_if(working_day_flag = 'Y')
        over(partition by country_alpha, trunc(dt_bord_ref, 'month') order by dt_bord_ref)
        month_to_date,
    count_if(working_day_flag = 'Y')
        over(partition by country_alpha, trunc(dt_bord_ref, 'month') order by dt_bord_ref
             rows between 1 following and unbounded following)
        month_to_go    
from df
""")

result.show()
+-------------------+-------------+----------------+-------------+-----------+
|        DT_BORD_REF|COUNTRY_ALPHA|working_day_flag|month_to_date|month_to_go|
+-------------------+-------------+----------------+-------------+-----------+
|2021-01-01 00:00:00|          BRA|               N|            0|          3|
|2021-01-02 00:00:00|          BRA|               N|            0|          3|
|2021-01-03 00:00:00|          BRA|               N|            0|          3|
|2021-01-04 00:00:00|          BRA|               Y|            1|          2|
|2021-01-05 00:00:00|          BRA|               Y|            2|          1|
|2021-01-06 00:00:00|          BRA|               Y|            3|          0|
|2021-01-01 00:00:00|          ITA|               N|            0|          3|
|2021-01-02 00:00:00|          ITA|               N|            0|          3|
|2021-01-03 00:00:00|          ITA|               N|            0|          3|
|2021-01-04 00:00:00|          ITA|               Y|            1|          2|
|2021-01-05 00:00:00|          ITA|               Y|            2|          1|
|2021-01-06 00:00:00|          ITA|               N|            2|          1|
|2021-01-07 00:00:00|          ITA|               Y|            3|          0|
|2021-01-01 00:00:00|          FRA|               N|            0|          3|
|2021-01-02 00:00:00|          FRA|               N|            0|          3|
|2021-01-03 00:00:00|          FRA|               N|            0|          3|
|2021-01-04 00:00:00|          FRA|               Y|            1|          2|
|2021-01-05 00:00:00|          FRA|               Y|            2|          1|
|2021-01-06 00:00:00|          FRA|               Y|            3|          0|
+-------------------+-------------+----------------+-------------+-----------+

如果您想在 Pyspark API 中使用类似的解决方案:

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

result = df.withColumn(
    'month_to_date',
    F.count(
        F.when(F.col('working_day_flag') == 'Y', 1)
    ).over(
        Window.partitionBy('country_alpha', F.trunc('dt_bord_ref', 'month'))
              .orderBy('dt_bord_ref')
    )
).withColumn(
    'month_to_go',
    F.count(
        F.when(F.col('working_day_flag') == 'Y', 1)
    ).over(
        Window.partitionBy('country_alpha', F.trunc('dt_bord_ref', 'month'))
              .orderBy('dt_bord_ref')
              .rowsBetween(1, Window.unboundedFollowing)
    )
)