UDAF Spark 中的多列输出

Question

我从我的 mongodb 中得到一些如下所示的数据：

     +------+-------+
     | view | data  |
     +------+-------+
     |  xx  | ***   |
     |  yy  | ***   |
     |  xx  | ***   |
     +------+-------+

其实没必要知道里面有什么。

我写了一个这样的 UserDefinedAggregateFunction 因为我想在视图上分组。：

class Extractor() extends UserDefinedAggregateFunction{
  override def inputSchema: StructType = // some stuff

  override def bufferSchema: StructType = 
      StructType(
        List(
          StructField("0",IntegerType,false),
          StructField("1",IntegerType,false),
          StructField("2",IntegerType,false),
          StructField("3",IntegerType,false),
          StructField("4",IntegerType,false),
          StructField("5",IntegerType,false),
          StructField("6",IntegerType,false),
          StructField("7",IntegerType,false)
        )
      )

  override def dataType: DataType = bufferSchema        

  override def deterministic: Boolean = true

  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    for (x <- 0 to 7){
      buffer(x) = 0
    }
  }

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = // some stuff

  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = // some stuff

  override def evaluate(buffer: Row): Any = 
      var l = List.empty[Integer]
      for (x <- 7 to 0 by -1){
          l = buffer.getInt(x) :: l
      }
      l
}

我的输出应该是这样的：

     +------+---+---+---+---+---+---+---+---+
     | view | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
     +------+---+---+---+---+---+---+---+---+
     |  xx  | 0 | 0 | 4 | 1 | 0 | 0 | 3 | 0 |
     |  yy  | 0 | 0 | 0 | 3 | 0 | 1 | 0 | 0 |
     +------+---+---+---+---+---+---+---+---+

这些值是在上面的更新/合并函数中计算出来的，但这是有效的，没有必要让你看到。

然后我这样使用它：

val ex = new Extractor()
val df = dataset.groupBy("view").agg(
      ex(dataset.col("data"))
)
df.show()

当我执行 df.show() 时，它总是给我一个 IndexOutOfBoundException。我知道这是懒惰的评估，这就是我在 df.show() 中收到错误的原因。

据我所知，它可以执行第一组并结束评估功能。但在那之后我得到一个 IndexOutOfBoundException...

当我更改数据类型并将函数评估为：

override def dataType: DataType =
    ArrayType(IntegerType,false)

override def evaluate(buffer: Row): Any = {
    var l = ofDim[Integer](8)
    for (x <- 0 to 7){
      l(x) = buffer.getInt(x)
    }
    l

输出如下所示：

     +------+------------------------------+
     | view | Extractor                    |
     +------+------------------------------+
     |  xx  | [0, 0, 4, 1, 0, 0, 3, 0]     |
     |  yy  | [0, 0, 0, 3, 0, 1, 0, 0]     |
     +------+------------------------------+

架构看起来像这样：

root
 |-- view: string (nullable = true)
 |-- Extractor: array (nullable = true)
 |    |-- element: integer (containsNull = false)

而且我无法将其转换为我想要的形式。

因为第二种方法有效，我想我在第一种方法中弄乱了 DataType 的东西，但我不知道如何解决它...

我的问题有很多介绍：

我怎样才能得到我想要的输出？ 我真的不在乎这两种方法中的哪一种（首先使用多个输出列或可以转换为我想要的形式的数组），只要它是有效的。

感谢您的帮助

Answer 1

您将聚合输出定义为列表：

 override def dataType: DataType = bufferSchema

因为bufferSchema 是一个List，所以这就是你最终得到的。您可以稍后更改架构并将列表中的每一列转换为新列。

对于你的错误，区别：

override def evaluate(buffer: Row): Any = 
  var l = List.empty[Integer]
  for (x <- 7 to 0 by -1){
      l = buffer.getInt(x) :: l
  }
  l

和

override def evaluate(buffer: Row): Any = 
  var l = ofDim[Integer](8)
  for (x <- 0 to 7){
    l = buffer.getInt(x) :: l
  }
  l

是在第二个中，您定义了预定义的列数。因此，您确信可以毫无问题地从 0 迭代到 7。

您的第一个示例并非如此，因此，我怀疑您的数据格式错误，导致您的缓冲区在initialize 或merge 中被错误地初始化。我建议您在转换缓冲区长度的每个步骤之后添加一个 try/catch 来验证大小（至少 initialize，但也可以是 update 或 merge）。

要为列表中的每个元素添加列，您可以使用 withColumn 或通过地图来完成。