pyspark 从星期四开始的一周获取月份的周数

Question

我需要获取一个月的工作日。但是一周的开始日应该是星期四。

数据如下

目前我的代码将开始日设为星期日。

df = df.withColumn("Week_Number",date_format(to_date("inv_dt", "yyyy-MM-dd"), "W"))

但是我希望一周从星期四开始

Answer 1

考虑创建一个用户定义函数 (UDF) 来处理特殊情况。

例如，以下 UDF 处理星期日以外的一周开始时的特殊情况。

参数start_day_of_week取一个整数作为 isoweekday ，其中星期一是 1，星期日是 7:

from pyspark.sql.functions import udf

def week_number(date_str, start_day_of_week=7):
    @udf("long")
    def _week_number(date_str):
        from datetime import datetime, date
        d = datetime.strptime(date_str, '%Y-%m-%d')     # compatible with Python2
        wd_d1 = date(d.year, d.month, 1).isoweekday()
        offset = (wd_d1 - start_day_of_week + 7 ) % 7
        return (d.day - 1 + offset) // 7 + 1
    return _week_number(date_str)

注意:上面的代码在 PySpark (2.4.4) 中使用 Python3 (3.7.3) 和 Python (2.7.15) 进行了测试。默认情况下，PySpark 从 Python (2.7.15) 开始。要使用 Python3，请使用 PYSPARK_PYTHON=python3 pyspark 启动 PySpark , 或添加 export PYSPARK_PYTHON=python3启动脚本。

如果 UDF 仅在 Python3 中使用，则可以简单地使用 d = date.fromisoformat(date_str)并且只导入 date来自 datetime .

演示

考虑如下输入数据:

from pyspark.sql.functions import *

inv_dt = [ '2018-09-{:02d}'.format(x) for x in range(1, 31) ]

# if with Pandas installed
import pandas as pd
df_pd = pd.DataFrame({'inv_dt':inv_dt})
df = spark.createDataFrame(df_pd)

# if without Pandas
from pyspark.sql.types import StringType
df = spark.createDataFrame(inv_dt, StringType()).withColumnRenamed('value', 'inv_dt')

df.show()

# +----------+
# |    inv_dt|
# +----------+
# |2018-09-01|
# |2018-09-02|
# |2018-09-03|
# |2018-09-04|
# |2018-09-05|
# |2018-09-06|
# |2018-09-07|
# |2018-09-08|
# |2018-09-09|
# |2018-09-10|
# |2018-09-11|
# |2018-09-12|
# |2018-09-13|
# |2018-09-14|
# |2018-09-15|
# |2018-09-16|
# |2018-09-17|
# |2018-09-18|
# |2018-09-19|
# |2018-09-20|
# +----------+
# only showing top 20 rows

并调用我们的 UDF week_number就像其他 Spark SQL 函数一样，例如 date_format .我们可以比较我们的 UDF 输出 week_number("inv_dt", 7)内置 Spark date_format(to_date("inv_dt", "yyyy-MM-dd"), "W") .

( df    
    .withColumn("Day_of_Week", date_format(to_date("inv_dt", "yyyy-MM-dd"), "u"))
    .withColumn("Week_Number_Sun_ref", date_format(to_date("inv_dt", "yyyy-MM-dd"), "W"))
    .withColumn("Week_Number_Sun", week_number("inv_dt", 7))
    .withColumn("Day_of_Week_Thu", week_number("inv_dt", 4))
    .withColumn("Day_of_Week_Mon", week_number("inv_dt", 1))
).show()

# +----------+-----------+-------------------+---------------+---------------+---------------+
# |    inv_dt|Day_of_Week|Week_Number_Sun_ref|Week_Number_Sun|Day_of_Week_Thu|Day_of_Week_Mon|
# +----------+-----------+-------------------+---------------+---------------+---------------+
# |2018-09-01|          6|                  1|              1|              1|              1|
# |2018-09-02|          7|                  2|              2|              1|              1|
# |2018-09-03|          1|                  2|              2|              1|              2|
# |2018-09-04|          2|                  2|              2|              1|              2|
# |2018-09-05|          3|                  2|              2|              1|              2|
# |2018-09-06|          4|                  2|              2|              2|              2|
# |2018-09-07|          5|                  2|              2|              2|              2|
# |2018-09-08|          6|                  2|              2|              2|              2|
# |2018-09-09|          7|                  3|              3|              2|              2|
# |2018-09-10|          1|                  3|              3|              2|              3|
# |2018-09-11|          2|                  3|              3|              2|              3|
# |2018-09-12|          3|                  3|              3|              2|              3|
# |2018-09-13|          4|                  3|              3|              3|              3|
# |2018-09-14|          5|                  3|              3|              3|              3|
# |2018-09-15|          6|                  3|              3|              3|              3|
# |2018-09-16|          7|                  4|              4|              3|              3|
# |2018-09-17|          1|                  4|              4|              3|              4|
# |2018-09-18|          2|                  4|              4|              3|              4|
# |2018-09-19|          3|                  4|              4|              3|              4|
# |2018-09-20|          4|                  4|              4|              4|              4|
# +----------+-----------+-------------------+---------------+---------------+---------------+
# only showing top 20 rows
#