如何在 Spark 中编写 UDF 以将索引映射到字符串标签？

Question

我正在使用 Spark，并且在名为 predictions 的列之一中有一个具有特定字符串格式的表。格式始终为 - 0=some_probability,1=some_other_probability,2=some_other_probability 类型。

以下是该表中的一些示例记录 -

val table1 = Seq(
("0=0.5,1=0.3,2=0.2"),
("0=0.6,1=0.2,2=0.2"),
("0=0.1,1=0.1,2=0.8")
).toDF("predictions")

table1.show(false)
+-----------------+
|predictions      |
+-----------------+
|0=0.5,1=0.3,2=0.2|
|0=0.6,1=0.2,2=0.2|
|0=0.1,1=0.1,2=0.8|
+-----------------+

现在，我还有关于每个索引的元数据信息 - 0,1,2...n 在单独的字符串中。元数据字符串看起来像 -

val metadata = "AA::BB::CC"

我想在 Scala 中编写一个 UDF 来将这些索引映射到字符串中的每个元素。该 UDF 的输出应该给我一个看起来像这样的新列 -

+--------------------+
|labelled_predictions|      
+--------------------+
|AA=0.5,BB=0.3,CC=0.2|
|AA=0.6,BB=0.2,CC=0.2|
|AA=0.1,BB=0.1,CC=0.8|
+--------------------+

所以，0 被AA 替换，因为AA 是metadata 字符串中始终被:: 分割的第一个元素。

如何在 Scala-Spark 中编写 UDF 来做到这一点？

Answer 1

val metadata = "AA::BB::CC"

根据给定的数据，这应该适合您：

def myUDF(metadata:String) = udf((s: String) => {
  val metadataSplit = metadata.split("::")
  val dataSplit = s.split(",")
  val output = new Array[String](dataSplit.size)
  for (i <- 0 until dataSplit.size) {
    output(i) = metadataSplit(i) + "=" + dataSplit(i).split("=")(1)
  }
  output.mkString(",")
})

table1.withColumn("labelled_predictions", myUDF(metadata)(col("predictions"))).select("labelled_predictions").show(false)

输出：

+--------------------+
|labelled_predictions|
+--------------------+
|AA=0.5,BB=0.3,CC=0.2|
|AA=0.6,BB=0.2,CC=0.2|
|AA=0.1,BB=0.1,CC=0.8|
+--------------------+