Question

我有一个如下数据框：

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

我需要创建一个新的数据帧，如果第一优先级为真（优先级priority1，priority2，priority3的顺序），我需要将所有其他优先级设置为false。寻找动态解决方案，我可以添加更多优先级列。

上述示例的预期输出：

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

Answer 1

首先，将所有优先级列放在一个数组中，并使用UDF将其转换为包含单个true值。然后，要将数组值放入各自的列中，请使用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变量中。这可以通过获取所有列名并应用过滤器来完成（如上例所示）列具有相似的名称。

Spark中的列式操作

1 个答案: