如果只使用SQL的解决方案是可以接受的，那么你可以使用窗口函数来生成标志列。一个例子：

with cte as (
    select *
         , sum(is_working_day) over (partition by datepart(year, day), datepart(month, day) order by day) as rn
         , sum(is_working_day) over (partition by datepart(year, day), datepart(month, day)) as mv
    from t
)
select day
     , is_working_day
     , case
           when is_working_day = 1 and rn >= 1 and rn <= 2 then rn
           when is_working_day = 1 and rn > mv - 3 then 5 - (mv - rn)
       end as flag
from cte

字符串
测试结果：
| 天|是工作日|标志|
| --|--|--|
| 2023-01-01 2023-01-01| 0 |* 空 |
| 2023-01-02 2023-01-02| 1 | 1 |
| 2023-01-03 2023-01-03| 0 | 空 |
| 2023-01-04 2023-01-04| 1 | 2 |
| 2023-01-05 2023-01-05| 1 | 3 |
| 2023-01-29 2023-01-29 2023-01-29| 1 | 4 |
| 2023-01-30 - 2023-01-30| 0 | 空 *|
| 2023-01-31 -01- 01 - 01 - 01| 1 | 5 |
| 2023-02-01 2023-02-01 2023-02-01| 1 | 1 |
| 2023-02-02 2023 -02-02| 1 | 2 |
将日期部分提取函数替换为RDBMS中可用的函数。
DB<>Fiddle

试试这个：

from pyspark.sql.functions import *
from pyspark.sql import Window
ab=ab.withColumn("Day",to_date(col("Day"),"dd-mm-yyyy"))
ab=ab.withColumn("Month",month(col("Day")))
w0=Window.partitionBy('month','is_working_day').orderBy('Day')
w1=Window.partitionBy('month','is_working_day')
ab=ab.withColumn("row",row_number().over(w0))


ab=ab.withColumn("flag",when((col("is_working_day")==1) & (col("row")==1),1)\
                       .when((col("is_working_day")==1) & (col("row")==2),2)\
                       .when((col("is_working_day")==1) & (col("row")==max(col("row")).over(w1)),5)\
                       .when((col("is_working_day")==1) & (col("row")==max(col("row")).over(w1)-1),4)\
                       .when((col("is_working_day")==1) & (col("row")==max(col("row")).over(w1)-2),3)\
                       .otherwise(0))

字符串