使用 scala 使用布尔运算折叠火花数据框中的列

Posted 2023-04-15

【中文标题】使用 scala 使用布尔运算折叠火花数据框中的列

【英文标题】：Collapse column in a spark dataframe using boolean operation using scala

【发布时间】：2020-06-13 09:58:07

【问题描述】：

我们如何使用 scala 使用 OR 操作将布尔列折叠为单行？ 第 1 部分：

A  true
A  false
B  false
B  false
C  true
B  false
A  true
C  true

期望的输出

B  false
A  true
C  true

我能想到的一个解决方案是按第一列条目对它们进行分组，在单独的数据框中过滤真假行，删除重复项，最后在检查字母 (例如 A) 已经存在于另一个（真实的）数据框中。

这个解决方案相当混乱。另外，不知道这是否适用于所有边缘情况。有什么聪明的方法可以做到这一点。

我是一个绝对的初学者，感谢任何帮助。 编辑：给定的答案适用于上述场景，但不适用于这种场景。有什么方法可以达到预期的输出？ 第 2 部分：

A  true    "Apple"
A  false   ""
B  false   ""
B  false   ""
C  true    "Cat"
C  true    "Cotton"
C  false   ""

期望的输出

B  false []
A  true  ["Apple"]
C  true  ["Cat","Cotton"]

我尝试通过按 col1 和 col2 分组然后使用 collect_set 折叠 col3 来实现这一点，然后

按第一列分组 收集第二列作为布尔值集 检查是否有一个 true 如果是，那么您的 OR 表达式将始终计算为 true。

但这会导致 col3_set 一起丢失。

【参考方案1】：

按第一列分组 收集第二列作为布尔值集 收集第三列作为字符串集 检查是否有一个 true，如果是，那么您的 OR 表达式将始终为真。 从col3_set 中删除空字符串""

import org.apache.spark.sql.functions._

object GroupByAgg 

  def main(args: Array[String]): Unit = 

    val spark = Constant.getSparkSess

    import spark.implicits._

    val df = List(("A", true,"Apple"), ("A", false,""),
      ("B", false,""),
      ("B", false,""),
      ("C", true,"Cat"),
      ("C", true,"Cotton"),
      ("C", true,"")).toDF("Col1", "Col2","Col3")

    //Group by 1st column
    df.groupBy("Col1")
      // Collect unique values
      .agg(collect_set("Col2").as("Col2_set"),collect_set("Col3").as("Col3_set"))
      //check if the array contains single true
      .withColumn("OutputCol2", when(array_contains(col("Col2_set"), true), true)
        .otherwise(false))
    .withColumn("OutputCol3",array_remove(col("Col3_set"),lit("")))
//.withColumn("OutputCol3",expr("filter(Col3_set, x -> x != '')"))
      .drop("Col2_set")
      .drop("Col3_set")
      .show()
  

输出：

+----+----------+-------------+
|Col1|OutputCol2|   OutputCol3|
+----+----------+-------------+
|   B|     false|           []|
|   C|      true|[Cat, Cotton]|
|   A|      true|      [Apple]|
+----+----------+-------------+

【讨论】：

我已经编辑了这个问题，这适用于第 1 部分，我怎样才能达到第 2 部分所需的输出。

@Nikita 更新了答案。您应该发布一个新问题而不是更新旧问题，因为审阅者可能会否决旧答案，因为他们与预期输出不匹配。

不知道，新来的，会小心的。感谢您更新答案:)

有没有办法在不使用 array_remove() 的情况下实现这一点

@Nikita 添加替代解决方案。 .withColumn("OutputCol3",expr("filter(Col3_set, x -> x != '')")) 在答案中

【参考方案2】：

试试这个：

按 col1 分组 collect_set on col2 -- 如果是布尔值，collect_set 会给你 最多 2 个元素，这对性能也有好处

将步骤 2 中收集的一组布尔值传递给 UDF，该 UDF 将执行简单的 recudeLeft 以对所有元素执行 OR 操作。

scala> val df = List(
     | ("A",  true),
     | ("A",  false),
     | ("B",  false),
     | ("B",  false),
     | ("C",  true),
     | ("B",  false),
     | ("A",  true),
     | ("C",  true)
     | ).toDF("col1","col2")
df: org.apache.spark.sql.DataFrame = [col1: string, col2: boolean]

scala> df.show
+----+-----+
|col1| col2|
+----+-----+
|   A| true|
|   A|false|
|   B|false|
|   B|false|
|   C| true|
|   B|false|
|   A| true|
|   C| true|
+----+-----+


scala> val aggOr = udf((a:Seq[Boolean])=>a.reduceLeft(_||_))
aggOr: org.apache.spark.sql.expressions.UserDefinedFunction = UserDefinedFunction(<function1>,BooleanType,Some(List(ArrayType(BooleanType,false))))

scala> df.groupBy("col1").agg(aggOr(collect_set("col2")).as("col2Or")).show
+----+------+
|col1|col2Or|
+----+------+
|   B| false|
|   C|  true|
|   A|  true|
+----+------+