Выберите литерал на основе значения столбца в Spark

Question

У меня есть карта:

val map = Map("A" -> 1, "B" -> 2)

И у меня есть DataFrame. столбец во фрейме данных содержит ключи на карте. Я пытаюсь выбрать столбец в новом DF, в котором есть значения карты на основе ключа:

val newDF = DfThatContainsTheKeyColumn.select(concat(col(SomeColumn), lit("|"),
    lit(map.get(col(ColumnWithKey).toString()).get) as newColumn)

Но это приводит к следующей ошибке:

java.lang.RuntimeException: Unsupported literal type class scala.None$ None

Я убедился, что столбец ColumnWithKey имеет только A s и B s и не содержит пустых значений.

Есть ли другой способ получить искомый результат? Любая помощь будет оценена.

Answer 1

Проблема в этом утверждении (помимо проблем с синтаксисом)

val newDF = DfThatContainsTheKeyColumn.select(concat(col(SomeColumn), lit("|"),
    lit(map.get(col(ColumnWithKey).toString()).get) as newColumn)

заключается в том, что col(ColumnWithKey) не будет принимать значение определенной строки c, а задано только схемой, то есть имеет постоянное значение.

В вашем случае я бы предложил присоединить вашу карту к вашему фрейму данных:

val map = Map("A" -> 1, "B" -> 2)
val df_map = map.toSeq.toDF("key","value")

val DfThatContainsTheKeyColumn = Seq(
  "A",
  "A",
  "B",
  "B"
).toDF("myCol")


DfThatContainsTheKeyColumn
  .join(broadcast(df_map),$"mycol"===$"key")
  .select(concat($"mycol",lit("|"),$"value").as("newColumn"))
  .show()

дает

|newColumn|
+---------+
|      A|1|
|      A|1|
|      B|2|
|      B|2|
+---------+

Answer 2

Вы можете использовать case classes, чтобы сделать это легко. Это пример:

С учетом этого ввода

val givenMap = Map("A" -> 1, "B" -> 2)

import spark.implicits._
val df = Seq(
  (1, "A"),
  (2, "A"),
  (3, "B"),
  (4, "B")
).toDF("col_a", "col_b")
df.show()

Вышеприведенный код выглядит так:

+-----+-----+
|col_a|col_b|
+-----+-----+
|    1|    A|
|    2|    A|
|    3|    B|
|    4|    B|
+-----+-----+

givenMap: scala.collection.immutable.Map[String,Int] = Map(A -> 1, B -> 2)
import spark.implicits._
df: org.apache.spark.sql.DataFrame = [col_a: int, col_b: string]

Код, который вам нужен, будет выглядеть так:

case class MyInput(col_a: Int, col_b: String)
case class MyOutput(col_a: Int, col_b: String, new_column: Int)

df.as[MyInput].map(row=> MyOutput(row.col_a, row.col_b, givenMap(row.col_b))).show()

С помощью case classes вы можете разыграть свой df и использовать объектную нотацию для доступа к значениям вашего столбца в пределах .map. Выше код будет выводить:

+-----+-----+----------+
|col_a|col_b|new_column|
+-----+-----+----------+
|    1|    A|         1|
|    2|    A|         1|
|    3|    B|         2|
|    4|    B|         2|
+-----+-----+----------+

defined class MyInput
defined class MyOutput

Answer 3

Вы можете искать карту, используя ключ из столбца, как,

val map = Map("A" -> 1, "B" -> 2)    
val df = spark.createDataset(Seq("dummy"))
      .withColumn("key",lit("A"))

df.map{ row =>
          val k = row.getAs[String]("key")
          val v = map.getOrElse(k,0)
        (k,v)
        }.toDF("key", "value").show(false)

Результат -

+---+-----+
|key|value|
+---+-----+
|A  |1    |
+---+-----+

Вы можете посмотреть карту, присутствующую внутри столбца, используя буквенный ключ, используя Column.getItem, см. Пример ниже.

    val mapKeys = Array("A","B")
    val mapValues = Array(1,2)

    val df = spark.createDataset(Seq("dummy"))
      .withColumn("key",lit("A"))
      .withColumn("keys",lit(mapKeys))
      .withColumn("values",lit(mapValues))
      .withColumn("map",map_from_arrays($"keys",$"values"))
      .withColumn("lookUpTheMap",$"map".getItem("A"))

      //A dataframe with Map is created.
      //A map is looked up using a hard coded String key.
      df.show(false)

Результат

+-----+---+------+------+----------------+------------+
|value|key|keys  |values|map             |lookUpTheMap|
+-----+---+------+------+----------------+------------+
|dummy|A  |[A, B]|[1, 2]|[A -> 1, B -> 2]|1           |
+-----+---+------+------+----------------+------------+

Чтобы найти карту внутри столбца на основе другого столбца, содержащего ключ - вы можете использовать UDF или используйте функцию отображения на кадре данных, как показано ниже.

    //A map is looked up using a Column key.
      df.map{ row =>
          val m = row.getAs[Map[String,Int]]("map")
          val k = row.getAs[String]("key")
          val v = m.getOrElse(k,0)
        (m,k,v)
        }.toDF("map","key", "value").show(false)

Результат

+----------------+---+-----+
|map             |key|value|
+----------------+---+-----+
|[A -> 1, B -> 2]|A  |1    |
+----------------+---+-----+