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

Question

假设我有一个如下所示的数据框

Id	A	B	C	D
1	100	10	20	5
2	0	5	10	5
3	0	7	2	3
4	0	1	3	7

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

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

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

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 列排序的窗口上使用collect_list 函数，并获取包含A 和max(B, C) + D 值的结构的累积数组（作为字段T）。然后，应用aggregate 计算列E。

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

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 said，您不能使用滞后函数来检索列的变化值。当您使用 scala API 时，您可以开发自己的 User-Defined Aggregate 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 转换为您在 input 数据帧上应用的用户定义聚合函数：

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")

如果您将问题的数据框作为input 数据框，您将获得以下result 数据框：

+---+---+---+---+---+---+
|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 |
+---+---+---+---+---+---+

Answer 3

我在 spark 2.3.0 中尝试过，但聚合函数有错误： 用户类抛出异常：Java.io.IOError: org.apache.spark.sql.catalyst.ParserException: externous input '>' excepting {'(','SELECT','FROM'....

我在本地尝试使用 spark 2.4.0 并且它可以工作，但在我们的 hdp 中我们有 2.3.0

你能帮我让它在 spark 2.3.0 上工作吗

提前致谢