PySpark计算物料风险周期Cut总和：窗口函数实现求助

需求拆解 #

针对每个Material：

1. 风险周期起始：第一个Cut非零值对应的Date
2. 风险周期结束：首次出现连续三个Cut=0之前的最后一个有效日期；若无连续三个零，则取最后一个Cut非零的日期；若全为零则无风险周期
3. Cut_Total：风险周期内所有Cut值的总和

实现代码 #

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

# 初始化SparkSession
spark = SparkSession.builder.appName("MaterialRiskCalculation").getOrCreate()

# ---------------------- 替换为你的数据加载逻辑 ----------------------
# 示例测试数据
sample_data = [
    ("A", "2023-01-01", 0),
    ("A", "2023-01-02", 5),
    ("A", "2023-01-03", 3),
    ("A", "2023-01-04", 0),
    ("A", "2023-01-05", 0),
    ("A", "2023-01-06", 0),
    ("A", "2023-01-07", 2),
    ("B", "2023-01-01", 2),
    ("B", "2023-01-02", 0),
    ("B", "2023-01-03", 4),
    ("B", "2023-01-04", 0),
    ("B", "2023-01-05", 0),
    ("C", "2023-01-01", 0),
    ("C", "2023-01-02", 0),
    ("C", "2023-01-03", 0),
    ("C", "2023-01-04", 1)
]
df = spark.createDataFrame(sample_data, ["Material", "Date", "Cut"])
df = df.withColumn("Date", F.to_date("Date"))
# -------------------------------------------------------------------

# 窗口定义：按物料分组，日期升序排序
w_material = Window.partitionBy("Material").orderBy("Date")

# 1. 标记非零Cut记录，确定风险起始日期
df = df.withColumn("is_non_zero", F.when(F.col("Cut") != 0, 1).otherwise(0))
df = df.withColumn("non_zero_cum", F.sum("is_non_zero").over(w_material))
df = df.withColumn("risk_start", F.first(F.when(F.col("non_zero_cum") == 1, F.col("Date"))).over(w_material))

# 2. 检测连续三个零的位置，确定风险结束日期
df = df.withColumn("prev1_non_zero", F.lag("is_non_zero", 1).over(w_material))
df = df.withColumn("prev2_non_zero", F.lag("is_non_zero", 2).over(w_material))
df = df.withColumn("is_three_consec_zero", F.when(
    (F.col("is_non_zero") == 0) & (F.col("prev1_non_zero") == 0) & (F.col("prev2_non_zero") == 0),
    1
).otherwise(0))

# 获取首次连续三个零的日期，结束日期为该日期前一天
df = df.withColumn("first_three_zero_date", F.first(F.when(F.col("is_three_consec_zero") == 1, F.col("Date"))).over(w_material))
df = df.withColumn("risk_end_candidate", F.date_sub(F.col("first_three_zero_date"), 1))

# 处理无连续三个零的情况：取最后一个非零日期，若无则为null
df = df.withColumn("last_non_zero_date", F.last(F.when(F.col("is_non_zero") == 1, F.col("Date"))).over(w_material))
df = df.withColumn("risk_end", F.coalesce(
    F.col("risk_end_candidate"),
    F.col("last_non_zero_date")
))

# 3. 计算风险周期内的Cut总和
df = df.withColumn("in_risk_period", F.when(
    (F.col("Date") >= F.col("risk_start")) & (F.col("Date") <= F.col("risk_end")),
    1
).otherwise(0))

result_df = df.groupBy("Material")\
    .agg(
        F.first("risk_start").alias("Risk_Start_Date"),
        F.first("risk_end").alias("Risk_End_Date"),
        F.sum(F.when(F.col("in_risk_period") == 1, F.col("Cut")).otherwise(0)).alias("Cut_Total")
    )\
    .orderBy("Material")

# 查看结果
result_df.show()

代码说明 #

• 风险起始确定：通过累积非零计数，定位每个物料第一条非零Cut的日期
• 连续零检测：使用lag函数获取前两条记录的非零标记，判断是否出现连续三个零；若出现，风险结束日期为该连续零的前一天
• 无连续零处理：若未出现连续三个零，风险结束日期取最后一个非零Cut的日期
• 总和计算：标记风险周期内的记录，分组求和得到Cut_Total

示例输出 #

+---------+---------------+-------------+---------+
|Material|Risk_Start_Date|Risk_End_Date|Cut_Total|
+---------+---------------+-------------+---------+
|        A|     2023-01-02|   2023-01-05|        8|
|        B|     2023-01-01|   2023-01-03|        6|
|        C|     2023-01-04|   2023-01-04|        1|
+---------+---------------+-------------+---------+

内容的提问来源于stack exchange，提问作者pyspark-lab