У меня есть фрейм данных, как показано ниже:

 Id  priority1 priority2 priority3
  1    true      true     true
  2    false     true     true
  3    false     false    false

Мне нужно создать новый фрейм данных, где, если первый приоритет истинен (порядок приоритета приоритет1, приоритет2, приоритет3), мне нужно сделать все другие приоритеты ложными. Ищу динамическое решение, в котором я могу добавить намного больше столбцов приоритета.

Ожидаемый результат из приведенного выше примера:

Id  priority1 priority2 priority3   new_priority1 new_priority2 new_priority3
1    true      true     true         true           false         false
2    false     true     true         false          true          false
3    false     false    true         false          false         true
1
John 14 Фев 2018 в 09:01

1 ответ

Лучший ответ

Сначала поместите все столбцы приоритета в массив и преобразуйте его, чтобы он содержал единственное истинное значение, используя UDF. Затем, чтобы поместить значения массива в их собственные столбцы, используйте foldLeft. Используя пример ввода:

val df = Seq((1, true, true, true), (2, false, true, true), (3, false, false, false))
  .toDF("Id", "priority1", "priority2", "priority3")

UDF и его использование:

val convertPriorities = udf((prios: Seq[Boolean]) => {
  val falseSeq = Seq.fill(prios.length)(false)
  prios.indexOf(true) match {
    case -1 => falseSeq
    case x => falseSeq.updated(x, true)
  }
})

val prioColumns = Seq("priority1", "priority2", "priority3")
val df2 = df.withColumn("priorities", convertPriorities(array(prioColumns.map(col(_)):_*)))

Обратите внимание, что переменная prioColumns создана, чтобы упростить foldLeft.

val df3 = prioColumns.zipWithIndex
  .foldLeft(df2)((df, col) => df.withColumn("new_" + col._1, $"priorities"(col._2)))
  .drop("priorities")

Это даст следующий результирующий фрейм данных:

+---+---------+---------+---------+-------------+-------------+-------------+
| Id|priority1|priority2|priority3|new_priority1|new_priority2|new_priority3|
+---+---------+---------+---------+-------------+-------------+-------------+
|  1|     true|     true|     true|         true|        false|        false|
|  2|    false|     true|     true|        false|         true|        false|
|  3|    false|    false|    false|        false|        false|        false|
+---+---------+---------+---------+-------------+-------------+-------------+

Используя этот подход, его должно быть очень легко расширить для использования большего количества столбцов, единственное изменение, которое необходимо сделать, - это добавить столбцы в переменную prioColumns. Это можно сделать, получив все имена столбцов и применив фильтр, если (как в примере выше) столбцы имеют похожие имена.

2
Shaido 14 Фев 2018 в 10:41