如何在Spark中基于前一行状态修改分组排序后DataFrame的行值

需求回顾 #

给定包含ID、DATE、State的DataFrame，需按以下规则更新每个ID分组内的State（按DATE排序后）：

• 当前行State为S且前一行State也为S时，改为FS
• 当前行State为E且前一行State也为E时，改为FE
• N保持不变

实现步骤（Python版本） #

1. 导入依赖

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lag, when
from pyspark.sql.window import Window

2. 创建/加载DataFrame

初始化SparkSession并处理示例数据（若已有DataFrame可跳过创建步骤）：

spark = SparkSession.builder.appName("StateAdjustment").getOrCreate()

# 示例数据
data = [
    ("X", "20-01-2023", "N"),
    ("X", "21-01-2023", "S"),
    ("X", "22-01-2023", "S"),
    ("X", "23-01-2023", "N"),
    ("X", "24-01-2023", "E"),
    ("X", "25-01-2023", "E"),
    ("Y", "20-01-2023", "S"),
    ("Y", "23-01-2023", "S")
]

df = spark.createDataFrame(data, ["ID", "DATE", "State"])
# 将字符串日期转为日期类型，确保排序逻辑正确
df = df.withColumn("DATE", col("DATE").cast("date", format="dd-MM-yyyy"))

3. 定义窗口规范

按ID分区、DATE升序排序，用于获取分组内前一行的State：

window_spec = Window.partitionBy("ID").orderBy("DATE")

4. 执行状态更新逻辑

用lag函数获取前一行状态，再通过条件判断修改当前行State：

result_df = df.withColumn("prev_state", lag(col("State")).over(window_spec)) \
    .withColumn(
        "State",
        when(
            (col("State") == "S") & (col("prev_state") == "S"),
            "FS"
        ).when(
            (col("State") == "E") & (col("prev_state") == "E"),
            "FE"
        ).otherwise(col("State"))
    ) \
    .drop("prev_state")  # 移除临时辅助列

5. 查看结果

result_df.orderBy("ID", "DATE").show()

输出结果与预期一致：

+---+----------+-----+
| ID|      DATE|State|
+---+----------+-----+
|  X|2023-01-20|    N|
|  X|2023-01-21|    S|
|  X|2023-01-22|   FS|
|  X|2023-01-23|    N|
|  X|2023-01-24|    E|
|  X|2023-01-25|   FE|
|  Y|2023-01-20|    S|
|  Y|2023-01-23|   FS|
+---+----------+-----+

实现步骤（Scala版本） #

若使用Scala开发，逻辑与Python一致：

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{col, lag, when}
import org.apache.spark.sql.expressions.Window

object StateAdjustment {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("StateAdjustment").getOrCreate()
    import spark.implicits._

    val data = Seq(
      ("X", "20-01-2023", "N"),
      ("X", "21-01-2023", "S"),
      ("X", "22-01-2023", "S"),
      ("X", "23-01-2023", "N"),
      ("X", "24-01-2023", "E"),
      ("X", "25-01-2023", "E"),
      ("Y", "20-01-2023", "S"),
      ("Y", "23-01-2023", "S")
    )

    val df = data.toDF("ID", "DATE", "State")
      .withColumn("DATE", col("DATE").cast("date"))

    val windowSpec = Window.partitionBy("ID").orderBy("DATE")

    val resultDf = df.withColumn("prev_state", lag(col("State")).over(windowSpec))
      .withColumn(
        "State",
        when(
          (col("State") === "S") && (col("prev_state") === "S"),
          "FS"
        ).when(
          (col("State") === "E") && (col("prev_state") === "E"),
          "FE"
        ).otherwise(col("State"))
      )
      .drop("prev_state")

    resultDf.orderBy("ID", "DATE").show()
  }
}

关键说明 #

• 窗口函数lag：用于获取分区内前一行的State值，默认偏移1行，可通过第二个参数调整偏移量
• 日期类型转换：避免字符串排序可能出现的逻辑错误（如"09-12-2023"与"10-01-2023"的字符串排序问题）
• 条件分支：嵌套when确保仅符合连续S/E的行被修改，其余状态保持不变

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