Scala 在Spark SQL中将数组作为UDF参数传递

我试图通过一个以数组为参数的函数来转换数据帧。我的代码如下所示：

def getCategory(categories:Array[String], input:String): String = { 
  categories(input.toInt) 
} 

val myArray = Array("a", "b", "c") 

val myCategories =udf(getCategory _ ) 

val df = sqlContext.parquetFile("myfile.parquet) 

val df1 = df.withColumn("newCategory", myCategories(lit(myArray), col("myInput")) 

但是，lit不喜欢数组，并且此脚本存在错误。我尝试定义一个新的部分应用函数，然后定义udf：

val newFunc = getCategory(myArray,  _:String) 
val myCategories = udf(newFunc) 

val df1 = df.withColumn("newCategory", myCategories(col("myInput"))) 

这也不起作用，因为我得到了一个nullPointer异常，并且myArray似乎没有被识别。关于如何将数组作为参数传递给具有数据帧的函数，有什么想法吗

---

另请注意，关于为什么在数据帧上使用函数这样简单的操作如此复杂（定义函数、将其重新定义为UDF等）的任何解释

很可能不是最漂亮的解决方案，但您可以尝试以下方法：

def getCategory(categories: Array[String]) = {
    udf((input:String) => categories(input.toInt))
}

df.withColumn("newCategory", getCategory(myArray)(col("myInput")))

您还可以尝试文本的

数组

：

val getCategory = udf(
   (input:String, categories: Array[String]) => categories(input.toInt))

df.withColumn(
  "newCategory", getCategory($"myInput", array(myArray.map(lit(_)): _*)))

在旁注中，使用

Map

而不是

Array

可能是一个更好的主意：

def mapCategory(categories: Map[String, String], default: String) = {
    udf((input:String) =>  categories.getOrElse(input, default))
}

val myMap = Map[String, String]("1" -> "a", "2" -> "b", "3" -> "c")

df.withColumn("newCategory", mapCategory(myMap, "foo")(col("myInput")))

自Spark 1.5.0以来，您还可以使用

阵列

功能：

import org.apache.spark.sql.functions.array

val colArray = array(myArray map(lit  _): _*)
myCategories(lit(colArray), col("myInput"))

---

另请参见

您现在可以使用typedLit。此外，您的输入应该是不可变的（Seq）。