spark dataframe udf将索引映射到值

时间:2017-04-13 17:57:12

标签: scala apache-spark dataframe indexing udf

我有一个火花数据框,其中一列由列表的索引组成。我想写一个udf,它允许我创建一个新的列,其中包含与索引相关的值。

E.g。

假设我有以下数据框和数组:

val df = spark.createDataFrame(Seq((0, Array(1, 1, 2)), (1, Array(1, 2, 0))))
df.show()
+---+---------+
| _1|       _2|
+---+---------+
|  0|[1, 1, 2]|
|  1|[1, 2, 0]|
+---+---------+
val sArray = Array("a", "b", "c")

我希望能够将_2中的指标映射到sArray中的值,从而导致这一点:

+---+---------+---------+
| _1|       _2|       _3|
+---+---------+---------+
|  0|[1, 1, 2]|[b, b, c]|
|  1|[1, 2, 0]|[b, c, a]|
+---+---------+---------+

我一直试图用udf:

来做这件事
def indexer (values: Array[String]) = 
  udf((indices: Array[Int]) => indices.map(values(_)))
df.withColumn("_3", indexer(sArray)($"_2"))

但是当我这样做时,我收到以下错误:

  

无法执行用户定义的功能

     

...引起:java.lang.ClassCastException:scala.collection.mutable.WrappedArray $ ofRef无法转换为[I

这里出了什么问题?我该如何解决这个问题?

1 个答案:

答案 0 :(得分:5)

在DataFrame中的ArrayType列上操作时,传递到UDF的实际类型为mutable.WrappedArray。您看到的失败是尝试将此WrappedArray投射到您的函数所期望的Array[Int]中的结果。

修复相当简单 - 定义期望mutable.WrappedArray[Int]

的函数
def indexer (values: Array[String]): UserDefinedFunction = {
  udf((indices: mutable.WrappedArray[Int]) => indices.map(values(_)))
}