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

发布时间：2020-12-16 08:56:27 所属栏目：安全来源：网络整理

导读：我试图通过一个以数组作为参数的函数来转换数据帧.我的代码看起来像这样： def getCategory(categories:Array[String],input:String): String = { categories(input.toInt) } val myArray = Array("a","b","c") val myCategories =udf(getCategory _ ) val d

我试图通过一个以数组作为参数的函数来转换数据帧.我的代码看起来像这样：

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

另见Spark UDF with varargs

（编辑：李大同）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!