PySpark实现按id1分组的季度成本汇总及同比增长率计算

此前代码的核心错误点 #

• 字符串截取逻辑错误：你写的substring(string(year_quarter), 2)存在两个问题：一是Spark SQL的字符串截取索引从1开始计数，和Python的0起始索引规则不同；二是没指定截取长度，会从起始位置一直取到字符串末尾。对于202001这类6位的year_quarter值，截取后会得到02001、02101这类完全不匹配的字符串，导致同季度的数据根本没被分到同一个窗口分区，每一行都是独立分区，自然拿不到上一年同季度的值，所有prev_value全为null。
• 窗口分区逻辑偏差：最初仅按id1分区排序时，lag函数默认取当前行的上一行，也就是相邻上一季度的值，并非同季度的上年值。
• 计算表达式笔误：差值计算里写的grouped_quarterly.cost - grouped_quarterly.cost是自身值相减，就算成功拿到prev_value，也算不出正确的差值结果。

可直接运行的完整代码 #

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

# 1. 按id1、季度聚合成本（修正原代码括号缺失问题）
grouped_quarterly = (
    df
    .withColumn(
        'year_quarter', 
        F.year(F.col('yyyy_mm_dd')) * 100 + F.quarter(F.col('yyyy_mm_dd'))
    )
    .groupby('id1', 'year_quarter')
    .agg(F.sum('cost').alias('cost'))
)

# 2. 定义同比计算的窗口规则
# 按id1、季度序号（year_quarter的后两位）分区，按时间排序
# 同分区内的行顺序为 2020Q1 → 2021Q1 → 2022Q1
yoy_window = Window.partitionBy(
    F.col('id1'),
    # 正确截取后两位的季度标识：从字符串第5位开始取2个字符，得到"01"/"02"/"03"/"04"
    F.substring(F.col('year_quarter').cast('string'), 5, 2)
).orderBy('year_quarter')

# 3. 计算同比增长率
growth_result = (
    grouped_quarterly
    # lag默认偏移1行，同分区内的上一行就是上一年同季度的成本值
    .withColumn('prev_year_cost', F.lag(F.col('cost')).over(yoy_window))
    # 按需求规则计算：上年同季度无数据则置0，否则计算百分比形式的增长率
    .withColumn(
        'cost_growth',
        F.when(
            F.col('prev_year_cost').isNull(),
            F.lit(0)
        ).otherwise(
            # 匹配样例的输出格式：取整数百分比值，如-81代表-81%
            F.round(
                (F.col('cost') - F.col('prev_year_cost')) / F.col('prev_year_cost') * 100, 
                0
            ).cast('int')
        )
    )
    # 选取最终需要的输出字段
    .select('year_quarter', 'id1', 'cost', 'cost_growth')
)

运行效果验证 #

用你给出的id1=222的测试数据运行后，输出结果为：

完全符合需求规则。

补充注意事项 #

• 计算除法时建议增加对prev_year_cost=0的异常判断，避免除零报错。
• 如果需要保留小数形式的增长率，只需调整round的精度参数即可。

内容的提问来源于stack exchange，提问作者stackq