spark dataframe udf 将索引映射到值

Question

我有一个 spark 数据框，其中一列由列表的索引组成。我想编写一个 udf，它允许我使用与索引关联的值创建一个新列。

例如

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

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

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

Answer 1

对 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(_)))
}