UDF scala return [最大，索引]

Question

我想为 Spark SQL 实现以下功能。给定一个数组 return 具有索引的最大值。我试过：

/*
 * 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))
}

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

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

返回一个 Array 没那么有用——因为索引类型总是 Int，而最大值类型取决于特定的调用（如果我理解正确，你希望它起作用适用于整数和双精度数）- 因此无法正确键入数组。

这是 UDF 的一种可能实现，返回一个元组 :

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

然后，可以调用 Doubles 以及 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]|
+-------------+

UDF scala return [最大，索引]

UDF scala return [max,index]

scala

user-defined-functions

apache-spark

apache-spark-sql