使用 Scala 根据前一行中不同列的计算值计算 Spark Dataframe 当前行中的列值

Question

假设我有一个 Dataframe 如下所示

ID	一个	乙	C	D
1	100	10	20	5
2	0	5	10	5
3	0	7	2	3
4	0	1	3	7

以上需要转换为类似下面的内容

ID	一个	乙	C	D	乙
1	100	10	20	5	75
2	75	5	10	5	60
3	60	7	2	3	50
4	50	1	3	7	40

这件事通过下面提供的细节起作用

1. 数据框现在有一个新列 E，第 1 行计算为col(A) - (max(col(B), col(C)) + col(D)) => 100-(max(10,20) + 5) = 75
2. 在Id为 2 的行中，来自第 1 行的 col E 的值作为 Col A的值向前推进
3. 因此，对于第 2 行，列E被确定为75-(max(5,10) + 5) = 60
4. 类似地，在Id为 3 的行中，A 的值变为 60，并且 col E的新值基于此确定

问题是，col A 的值取决于除第一行之外的前一行的值

是否有可能使用窗口和滞后来解决这个问题

Answer 1

您可以在按Id列排序的 Window 上使用collect_list function 并获取包含A和max(B, C) + D值的结构的累积数组（作为字段T ）。 然后，应用aggregate来计算列E 。

请注意，在这种特殊情况下，您不能使用lag window function 因为您希望递归地获取计算值。

import org.apache.spark.sql.expressions.Window

val df2 = df.withColumn(
  "tmp",
  collect_list(
    struct(col("A"), (greatest(col("B"), col("C")) + col("D")).as("T"))
  ).over(Window.orderBy("Id"))
).withColumn(
  "E",
  expr("aggregate(transform(tmp, (x, i) -> IF(i=0, x.A - x.T, -x.T)), 0, (acc, x) -> acc + x)")
).withColumn(
  "A",
  col("E") + greatest(col("B"), col("C")) + col("D")
).drop("tmp")

df2.show(false)

//+---+---+---+---+---+---+
//|Id |A  |B  |C  |D  |E  |
//+---+---+---+---+---+---+
//|1  |100|10 |20 |5  |75 |
//|2  |75 |5  |10 |5  |60 |
//|3  |60 |7  |2  |3  |50 |
//|4  |50 |1  |3  |7  |40 |
//+---+---+---+---+---+---+

您可以显示中间列tmp以了解计算背后的逻辑。

Answer 2

正如blackbishop 所说，您不能使用滞后 function 来检索列的变化值。 当您使用 scala API 时，您可以开发自己的用户定义聚合 Function

您创建以下案例类，代表您当前正在读取的行和聚合器的缓冲区：

case class InputRow(A: Integer, B: Integer, C: Integer, D: Integer)

case class Buffer(var E: Integer, var A: Integer)

然后你使用它们来定义你的RecursiveAggregator自定义聚合器：

import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.Encoder

object RecursiveAggregator extends Aggregator[InputRow, Buffer, Buffer] {
  override def zero: Buffer = Buffer(null, null)

  override def reduce(buffer: Buffer, currentRow: InputRow): Buffer = {
    buffer.A = if (buffer.E == null) currentRow.A else buffer.E
    buffer.E = buffer.A - (math.max(currentRow.B, currentRow.C) + currentRow.D)
    buffer
  }

  override def merge(b1: Buffer, b2: Buffer): Buffer = {
    throw new NotImplementedError("should be used only over ordered window")
  }

  override def finish(reduction: Buffer): Buffer = reduction

  override def bufferEncoder: Encoder[Buffer] = ExpressionEncoder[Buffer]

  override def outputEncoder: Encoder[Buffer] = ExpressionEncoder[Buffer]
}

最后，将RecursiveAggregator转换为用户定义的聚合 function，并将其应用于input dataframe：

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{col, udaf}

val recursiveAggregator = udaf(RecursiveAggregator)

val window = Window.orderBy("Id")

val result = input
  .withColumn("computed", recursiveAggregator(col("A"), col("B"), col("C"), col("D")).over(window))
  .select("Id", "computed.A", "B", "C", "D", "computed.E")

如果您将问题的 dataframe 作为input dataframe，您将得到以下result dataframe：

+---+---+---+---+---+---+
|id |A  |B  |C  |D  |E  |
+---+---+---+---+---+---+
|1  |100|10 |20 |5  |75 |
|2  |75 |5  |10 |5  |60 |
|3  |60 |7  |2  |3  |50 |
|4  |50 |1  |3  |7  |40 |
+---+---+---+---+---+---+