Как эффективно создать несколько столбцов из столбцов MapType (без смещения) - PullRequest
Новая
       6

Как эффективно создать несколько столбцов из столбцов MapType (без смещения)

0 голосов
/ 31 января 2020

Моя цель - создать столбцы из другого MapType столбца. Имена столбцов, являющихся ключами карты, и связанные с ними значения.

Ниже моего начального кадра данных: 

+-----------+---------------------------+
|id         |         mapColumn         |
+-----------+---------------------------+
| 1         |Map(keyA -> 0, keyB -> 1)  |
| 2         |Map(keyA -> 4, keyB -> 2)  |
+-----------+---------------------------+

Ниже желаемого результата: 

+-----------+----+----+
|id         |keyA|keyB|
+-----------+----+----+
| 1         |   0|   1|
| 2         |   4|   2|
+-----------+----+----+

Я нашел решение с помощью Foldleft с аккумуляторами (работают, но очень медленно): 

val colsToAdd = startDF.collect()(0)(1).asInstanceOf[Map[String,Integer]].map(x => x._1).toSeq
res1: Seq[String] = List(keyA, keyB)

val endDF = colsToAdd.foldLeft(startDF)((startDF, key) => startDF.withColumn(key, lit(0)))

//(lit(0) for testing)

Реальный стартовый фрейм данных огромен, мне нужна оптимизация.

1 Ответ

0 голосов
/ 31 января 2020

Вы можете просто использовать функцию explode, чтобы взорвать столбец типа карты, а затем использовать pivot, чтобы получить каждый ключ как новый столбец. Примерно так: 

val df = Seq((1,Map("keyA" -> 0, "keyB" -> 1)), (2,Map("keyA" -> 4, "keyB" -> 2))
).toDF("id", "mapColumn")

df.select($"id", explode($"mapColumn"))
  .groupBy($"id")
  .pivot($"key")
  .agg(first($"value"))
  .show()

Дает: 

+---+----+----+
| id|keyA|keyB|
+---+----+----+
|  1|   0|   1|
|  2|   4|   2|
+---+----+----+
...