scala - 过滤 Spark 数据帧中的多列

Question

假设我在 Spark 中有一个数据框，如下所示 -

val df = Seq(
(0,0,0,0.0),
(1,0,0,0.1),
(0,1,0,0.11),
(0,0,1,0.12),
(1,1,0,0.24),
(1,0,1,0.27),
(0,1,1,0.30),
(1,1,1,0.40)
).toDF("A","B","C","rate")

这是它的样子-

scala> df.show()
+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  0|  0|  0| 0.0|
|  1|  0|  0| 0.1|
|  0|  1|  0|0.11|
|  0|  0|  1|0.12|
|  1|  1|  0|0.24|
|  1|  0|  1|0.27|
|  0|  1|  1| 0.3|
|  1|  1|  1| 0.4|
+---+---+---+----+

A、B 和 C 是本例中的广告 channel 。 0 和 1 分别表示没有和存在 channel 。 2^3 显示数据框中的 8 种组合。

我想从此数据框中过滤一次显示 2 个 channel (AB、AC、BC)存在的记录。这是我希望输出的方式 -

+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  1|  1|  0|0.24|
|  1|  0|  1|0.27|
|  0|  1|  1| 0.3|
+---+---+---+----+

我可以写 3 条语句来获得输出 -

scala> df.filter($"A" === 1 && $"B" === 1 && $"C" === 0).show()
+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  1|  1|  0|0.24|
+---+---+---+----+


scala> df.filter($"A" === 1 && $"B" === 0  && $"C" === 1).show()
+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  1|  0|  1|0.27|
+---+---+---+----+


scala> df.filter($"A" === 0 && $"B" === 1 && $"C" === 1).show()
+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  0|  1|  1| 0.3|
+---+---+---+----+

但是，我想使用完成我的工作的单个语句或帮助我获取输出的函数来实现此目的。我正在考虑使用 case 语句来匹配值。但是一般来说，我的数据框可能包含 3 个以上的 channel -

scala> df.show()
+---+---+---+---+----+
|  A|  B|  C|  D|rate|
+---+---+---+---+----+
|  0|  0|  0|  0| 0.0|
|  0|  0|  0|  1| 0.1|
|  0|  0|  1|  0| 0.1|
|  0|  0|  1|  1|0.59|
|  0|  1|  0|  0| 0.1|
|  0|  1|  0|  1|0.89|
|  0|  1|  1|  0|0.39|
|  0|  1|  1|  1| 0.4|
|  1|  0|  0|  0| 0.0|
|  1|  0|  0|  1|0.99|
|  1|  0|  1|  0|0.49|
|  1|  0|  1|  1| 0.1|
|  1|  1|  0|  0|0.79|
|  1|  1|  0|  1| 0.1|
|  1|  1|  1|  0| 0.1|
|  1|  1|  1|  1| 0.1|
+---+---+---+---+----+

在这种情况下，我希望输出为 -

scala> df.show()
+---+---+---+---+----+
|  A|  B|  C|  D|rate|
+---+---+---+---+----+
|  0|  0|  1|  1|0.59|
|  0|  1|  0|  1|0.89|
|  0|  1|  1|  0|0.39|
|  1|  0|  0|  1|0.99|
|  1|  0|  1|  0|0.49|
|  1|  1|  0|  0|0.79|
+---+---+---+---+----+

显示 channel 成对存在的速率 =>(AB、AC、AD、BC、BD、CD)。

请帮忙。

Answer 1

一种方法是对列求和，然后仅当求和结果为 2 时才进行过滤。

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

df.withColumn("res", $"A" + $"B" + $"C").filter($"res" === lit(2)).drop("res").show

输出是:

+---+---+---+----+
|  A|  B|  C|rate|
+---+---+---+----+
|  1|  1|  0|0.24|
|  1|  0|  1|0.27|
|  0|  1|  1| 0.3|
+---+---+---+----+