Как написать UDF в Spark для отображения индексов на строковые метки?

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"

Я хотел бы написать UDF в Scala, чтобы сопоставить эти индексы с каждым элементом в строке. Выходные данные этого 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который всегда делится на ::.

Как мне написать UDF в Scala-Spark для этого?

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|
+--------------------+