UDF斯卡拉回报[最大，指数]

Question

我想实现以下函数来激发SQL。给定一个数组，返回索引的最大值。我试过：

/*
 * This function finds the maximum value and corresponding index in the array. NULLs are ignored. 
 * Return type is array in format [max, index], and its element type is the same as the input type.
 * Parameters: x Array[Int]
 * Returns: Array as [max, index].
 */
def array_max_index(x: WrappedArray[Int]): WrappedArray[Int] = {
    val arr = collection.mutable.WrappedArray.empty
    arr.:+(x.max).:+(x.indexOf(x.max))
}

这很好，但仅适用于整数-我希望UDF适用于其他数值（例如Doubles）。我尝试了以下方法，但我无法返回带有类型的结构：

def array_max_index[T](item:Traversable[T])(implicit n:Numeric[T]): Traversable[T] = {
    val arr = collection.mutable.WrappedArray.empty
    val max = item.max
    val index = n.toInt(item.toSeq.indexOf(max))
    arr.:+(max).:+(index)
  }

有什么想法吗？

Answer 1

返回＜code＞数组，并且最大值类型取决于特定调用（如果我理解正确，您希望它能很好地用于整数和双精度）-因此无法正确键入数组。

以下是UDF的一个可能实现，返回元组：

def array_max_index[T](x: Traversable[T])(implicit n: Numeric[T]): (T, Int) = {
  (x.max, x.toSeq.indexOf(x.max))
}

然后，可以为双s 和 Ints 调用：

sqlContext.udf.register("array_max_index", array_max_index(_: Traversable[Double]))

sqlContext.sql(
  """SELECT array_max_index(array(
    |  CAST(5.0 AS DOUBLE),
    |  CAST(7.0 AS DOUBLE),
    |  CAST(3.0 AS DOUBLE)
    |)) as max_and_index""".stripMargin).show

哪些打印：

+-------------+
|max_and_index|
+-------------+
|      [7.0,1]|
+-------------+