PySpark实现：按组计算当前行与组内首行的数值差值

没问题，我来帮你搞定这个需求——按组计算a、b列当前行与组内首行的数值差。下面是具体实现步骤和代码：

核心思路 #

我们需要借助**窗口函数（Window Functions）**来实现：

• 先按分组键（这里显然是date和t，示例中这两个字段的组内值一致）划分窗口
• 在每个窗口内，获取a、b列的首行值
• 用当前行的a、b分别减去对应首行值，得到差值列

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import first, col

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

# 模拟你的DataFrame数据（按示例截取部分）
data = [
    ("20180328", "1", "2018-W10", 31, 35),
    ("20180328", "1", "2018-W11", 18, 37),
    ("20180328", "1", "2018-W12", 19, 37),
    ("20180328", "1", "2018-W13", 19, 38),
    ("20180328", "1", "2018-W14", 20, 38),
    ("20180328", "1", "2018-W15", 22, 39),
    ("20180328", "1", "2018-W16", 23, 39),
    ("20180328", "1", "2018-W17", 24, 40),
    ("20180328", "1", "2018-W18", 25, 40)
]
df = spark.createDataFrame(data, ["date", "t", "week", "a", "b"])

# 定义窗口：按date和t分组，按week排序（保证首行是组内最早的周记录）
window_spec = Window.partitionBy("date", "t").orderBy("week")

# 添加差值列：当前行a/b 减去组内首行的a/b
result_df = df.withColumn("a_diff", col("a") - first("a").over(window_spec)) \
              .withColumn("b_diff", col("b") - first("b").over(window_spec))

# 查看结果
result_df.show()

代码解释 #

1. 窗口定义：partitionBy("date", "t")确保我们在同一date和t的组内计算；orderBy("week")保证组内首行是最早的周数据（如果你的业务不需要按week排序，可以去掉，但建议保留以确保结果的确定性）。
2. first函数：first("a").over(window_spec)会获取当前窗口内a列的第一个值，也就是组内首行的a值。
3. 差值计算：直接用当前行的a/b减去首行值，得到a_diff和b_diff列。

示例输出 #

运行代码后，你会得到类似这样的结果（截取前几行）：

+---------+---+--------+---+---+------+------+
|     date|  t|    week|  a|  b|a_diff|b_diff|
+---------+---+--------+---+---+------+------+
|20180328 |  1|2018-W10| 31| 35|     0|     0|
|20180328 |  1|2018-W11| 18| 37|   -13|     2|
|20180328 |  1|2018-W12| 19| 37|   -12|     2|
|20180328 |  1|2018-W13| 19| 38|   -12|     3|
+---------+---+--------+---+---+------+------+

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