Используйте карту для замены значений столбцов в Spark

Question

Мне нужно сопоставить список столбцов с другим столбцом в наборе данных Spark: подумайте как-то так

val translationMap: Map[Column, Column] = Map(
  lit("foo") -> lit("bar"),
  lit("baz") -> lit("bab")
)

И у меня есть такой фрейм данных:

val df = Seq("foo", "baz").toDF("mov")

Поэтому я намерен выполнить перевод следующим образом:

df.select(
  col("mov"),
  translationMap(col("mov"))
)

но этот фрагмент кода выдает следующую ошибку

key not found: movs
java.util.NoSuchElementException: key not found: movs

Есть ли способ выполнить такой перевод без объединения сотен when с? думаю, что translationMap может иметь много пар ключ-значение.

Answer 1

Вместо Map[Column, Column] вы должны использовать Column, содержащий литерал карты:

import org.apache.spark.sql.functions.typedLit

val translationMap: Column = typedLit(Map(
  "foo" -> "bar",
  "baz" -> "bab"
))

Остальная часть вашего кода может оставаться как есть:

df.select(
  col("mov"),
  translationMap(col("mov"))
).show

+---+---------------------------------------+
|mov|keys: [foo,baz], values: [bar,bab][mov]|
+---+---------------------------------------+
|foo|                                    bar|
|baz|                                    bab|
+---+---------------------------------------+

Answer 2

Вы не можете ссылаться на коллекцию Scala, объявленную в драйвере, как это, внутри распределенного фрейма данных. Альтернативой может быть использование UDF, которая не будет эффективной с точки зрения производительности, если у вас большой набор данных, поскольку UDF не оптимизируются Spark.

val translationMap = Map( "foo" -> "bar" , "baz" -> "bab" )
val getTranslationValue = udf ((x: String)=>translationMap.getOrElse(x,null.asInstanceOf[String]) )
df.select(col("mov"), getTranslationValue($"mov").as("value")  ).show

//+---+-----+
//|mov|value|
//+---+-----+
//|foo|  bar|
//|baz|  bab|
//+---+-----+

Другим решением было бы загрузить Map как DataSet[(String, String)] и объединить два набора данных, взяв mov в качестве ключа.