Spark: получить значение строки из окна

Question

С помощью spark я определил окно:

val window = Window
  .partitionBy("myaggcol")
  .orderBy("datefield")
  .rowsBetween(-2, 0)

Затем я могу вычислить новый столбец из строк окна, например:

dataset
  .withColumn("newcol", last("diffcol").over(window) - first("diffcol").over(window))

Это будет вычислять для каждой точкиразница в "diffcol" со строкой n-2.

Теперь мой вопрос: как я могу получить "diffcol" из строки n-1, не первый и не последний, а промежуточный?

Answer 1

Если я правильно понимаю ваш вопрос, оконная функция lag будет работать лучше, чем rowsBetween, как показано в следующем примере:

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

val df = Seq(
  ("a", 1, 100), ("a", 2, 200), ("a", 3, 300), ("a", 4, 400),
  ("b", 1, 500), ("b", 2, 600), ("b", 3, 700)
).toDF("c1", "c2", "c3")

val win = Window.partitionBy("c1").orderBy("c2")

df.
  withColumn("c3Diff1", $"c3" - coalesce(lag("c3", 1).over(win), lit(0))).
  withColumn("c3Diff2", $"c3" - coalesce(lag("c3", 2).over(win), lit(0))).
  show
// +---+---+---+-------+-------+
// | c1| c2| c3|c3Diff1|c3Diff2|
// +---+---+---+-------+-------+
// |  b|  1|500|    500|    500|
// |  b|  2|600|    100|    600|
// |  b|  3|700|    100|    200|
// |  a|  1|100|    100|    100|
// |  a|  2|200|    100|    200|
// |  a|  3|300|    100|    200|
// |  a|  4|400|    100|    200|
// +---+---+---+-------+-------+