如何用Scala DataFrame处理重复ID并更新对应end_date列值

现有如下示例数据：

由于id为2222的记录存在重复，需将PQR对应行的end_date更新为重复记录（XYZ）的start_date，请求用Scala DataFrame实现该需求。

可以通过窗口函数分组获取同id下的后续记录start_date，再结合条件更新实现需求，具体代码如下：

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

object UpdateEndDate {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("UpdateEndDate")
      .master("local[*]")
      .getOrCreate()

    import spark.implicits._

    // 构造示例数据
    val df = Seq(
      (1111, "ABC", "2020-01-01", ""),
      (2222, "PQR", "2020-01-01", "2020-06-30"),
      (2222, "XYZ", "2020-06-30", "")
    ).toDF("id", "name", "start_date", "end_date")

    // 定义窗口：按id分组，按start_date升序排列
    val windowSpec = Window.partitionBy("id").orderBy("start_date")

    // 获取同id下下一条记录的start_date
    val dfWithNextStart = df.withColumn("next_start_date", lead("start_date", 1).over(windowSpec))

    // 条件更新end_date：当name为PQR且存在next_start_date时，替换为next_start_date，否则保留原end_date
    val resultDf = dfWithNextStart.withColumn("end_date",
      when(col("name") === "PQR" && col("next_start_date").isNotNull, col("next_start_date"))
        .otherwise(col("end_date"))
    ).drop("next_start_date")

    // 展示结果
    resultDf.show()

    spark.stop()
  }
}

代码说明 #

• 窗口函数lead：用于获取同id分组内，当前行的下一行记录的start_date，作为后续更新的目标值。
• 条件更新when：针对name为PQR的行，若存在下一条记录的start_date，则将其end_date替换为该值，其他行保持原end_date不变。
• 最后删除辅助列next_start_date，得到最终结果。

执行结果 #

+----+----+----------+----------+
|  id|name|start_date|  end_date|
+----+----+----------+----------+
|1111| ABC|2020-01-01|          |
|2222| PQR|2020-01-01|2020-06-30|
|2222| XYZ|2020-06-30|          |
+----+----+----------+----------+

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