如何在Spark中实现列值向下滚动（Roll）操作？

问题分析 #

你需要的是循环移位：每一行的lag值取上一行的sample_column，第一行的lag值替换为当前分区（或全局）的最后一个sample_column值。之前的错误是因为lag函数的第三个参数必须是字面量（常量值），不能直接传入窗口函数返回的Column类型结果。

解决方案 #

Spark没有直接提供循环移位的内置函数，但可以通过组合现有窗口函数实现，以下是两种常用方案：

方案1：窗口函数+Coalesce替换NULL

无需额外关联，直接在窗口中获取分区最后值，再替换lag产生的NULL：

// 定义原分区排序窗口
val windowSpec = Window.partitionBy("file_name").orderBy("sample_column")
// 定义用于获取分区全局最后值的窗口（需和原窗口排序一致，确保取到排序后的最后元素）
val partitionLastWindow = Window.partitionBy("file_name").orderBy("sample_column")
  .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

val resultDF = df
  // 计算普通lag值，第一行为NULL
  .withColumn("lag", lag("sample_column", 1).over(windowSpec))
  // 获取当前分区的最后一个sample_column值
  .withColumn("partition_last", last("sample_column").over(partitionLastWindow))
  // 用分区最后值替换NULL
  .withColumn("lag", coalesce($"lag", $"partition_last"))
  .drop("partition_last")

resultDF.show()

方案2：分组聚合+关联替换

先通过分组获取每个分区的最后值，再关联回原表替换NULL：

val windowSpec = Window.partitionBy("file_name").orderBy("sample_column")

// 第一步：计算每个分区的最后一个sample_column值
val partitionLastDF = df
  .groupBy("file_name")
  .agg(last("sample_column").as("partition_last"))

// 第二步：计算lag并关联替换NULL
val resultDF = df
  .withColumn("lag", lag("sample_column", 1).over(windowSpec))
  .join(partitionLastDF, Seq("file_name"), "left")
  .withColumn("lag", coalesce($"lag", $"partition_last"))
  .drop("partition_last")

resultDF.show()

无分区场景简化版 #

如果不需要按分区处理，直接对整个DataFrame操作：

val windowSpec = Window.orderBy("sample_column")
val globalLastWindow = Window.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

val resultDF = df
  .withColumn("lag", lag("sample_column", 1).over(windowSpec))
  .withColumn("global_last", last("sample_column").over(globalLastWindow))
  .withColumn("lag", coalesce($"lag", $"global_last"))
  .drop("global_last")

resultDF.show()

内容的提问来源于stack exchange，提问作者Ícaro Lorran