使用 Scala 在 Spark 数据帧中加入数据

Question

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

val df = Seq(
(0,0,0,0.0,0),
(1,0,0,0.1,1),
(0,1,0,0.11,1),
(0,0,1,0.12,1),
(1,1,0,0.24,2),
(1,0,1,0.27,2),
(0,1,1,0.3,2),
(1,1,1,0.4,3)
).toDF("A","B","C","rate","total")

这是它的样子

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

在这种情况下，A、B 和 C 是通道。 0 和 1 分别代表通道的不存在和存在。 2^3 显示了数据帧中的 8 种组合，其中“总计”列给出了这 3 个通道的逐行总和。

这些通道出现的个别概率可以由 -

scala> val oneChannelCase = df.filter($"total" === 1).toDF()

scala> oneChannelCase.show()
+---+---+---+----+-----+
|  A|  B|  C|rate|total|
+---+---+---+----+-----+
|  1|  0|  0| 0.1|    1|
|  0|  1|  0|0.11|    1|
|  0|  0|  1|0.12|    1|
+---+---+---+----+-----+

但是，我只对这些通道的成对概率感兴趣 -

scala> val probs = df.filter($"total" === 2).toDF()

scala> probs.show()
+---+---+---+----+-----+
|  A|  B|  C|rate|total|
+---+---+---+----+-----+
|  1|  1|  0|0.24|    2|
|  1|  0|  1|0.27|    2|
|  0|  1|  1| 0.3|    2|
+---+---+---+----+-----+

我想做的是 - 将 3 个新列附加到这些“probs”数据框中，以显示各个概率。以下是我正在寻找的输出 -

A   B   C   rate    prob_A   prob_B   prob_C
1   1   0   0.24      0.1      0.11      0
1   0   1   0.27      0.1      0         0.12                     
0   1   1   0.3       0        0.11      0.12 

为了清楚起见，输出结果的第一行显示 A=1, B=1, C=0。因此，A=0.1、B=0.11 和 C=0 的个别概率分别附加到 probs 数据帧。同样，对于第二行，A=1、B=0、C=1 显示了 A=0.1、B=0 和 C=0.12 的单独概率分别附加到 probs 数据帧。

这是我尝试过的 -

scala> val channels = df.columns.filter(v => !(v.contains("rate") |  v.contains("total")))
#channels: Array[String] = Array(A, B, C)


scala> val pivotedProb = channels.map(v => f"case when $v = 1 then rate else 0 end as prob_${v}")

scala> val param = pivotedProb.mkString(",")

scala> val probs = spark.sql(f"select *, $param from df")

scala> probs.show()
+---+---+---+----+-----+------+------+------+
|  A|  B|  C|rate|total|prob_A|prob_B|prob_C|
+---+---+---+----+-----+------+------+------+
|  0|  0|  0| 0.0|    0|   0.0|   0.0|   0.0|
|  1|  0|  0| 0.1|    1|   0.1|   0.0|   0.0|
|  0|  1|  0|0.11|    1|   0.0|  0.11|   0.0|
|  0|  0|  1|0.12|    1|   0.0|   0.0|  0.12|
|  1|  1|  0|0.24|    2|  0.24|  0.24|   0.0|
|  1|  0|  1|0.27|    2|  0.27|   0.0|  0.27|
|  0|  1|  1| 0.3|    2|   0.0|   0.3|   0.3|
|  1|  1|  1| 0.4|    3|   0.4|   0.4|   0.4|
+---+---+---+----+-----+------+------+------+

这给了我错误的输出。

请帮忙。

Answer 1

如果我正确理解您的要求，使用 foldLeft 遍历通道列，您可以 1) 从单通道数据帧生成 ratesMap，以及 2) 将列添加到具有列的双通道数据帧值等于通道乘积和对应的ratesMap 值：

val df = Seq(
  (0, 0, 0, 0.0, 0),
  (1, 0, 0, 0.1, 1),
  (0, 1, 0, 0.11, 1),
  (0, 0, 1, 0.12, 1),
  (1, 1, 0, 0.24, 2),
  (1, 0, 1, 0.27, 2),
  (0, 1, 1, 0.3, 2),
  (1, 1, 1, 0.4, 3)
).toDF("A", "B", "C", "rate", "total")

val oneChannelDF = df.filter($"total" === 1)
val twoChannelDF = df.filter($"total" === 2)

val channels = df.columns.filter(v => !(v.contains("rate") || v.contains("total")))
// channels: Array[String] = Array(A, B, C)

val ratesMap = channels.foldLeft( Map[String, Double]() ){ (acc, c) =>
  acc + (c -> oneChannelDF.select("rate").where(col(c) === 1).head.getDouble(0))
}
// ratesMap: scala.collection.immutable.Map[String,Double] = Map(A -> 0.1, B -> 0.11, C -> 0.12)

val probsDF = channels.foldLeft( twoChannelDF ){ (acc, c) =>
  acc.withColumn( "prob_" + c, col(c) * ratesMap.getOrElse(c, 0.0) )
}

probsDF.show
// +---+---+---+----+-----+------+------+------+
// |  A|  B|  C|rate|total|prob_A|prob_B|prob_C|
// +---+---+---+----+-----+------+------+------+
// |  1|  1|  0|0.24|    2|   0.1|  0.11|   0.0|
// |  1|  0|  1|0.27|    2|   0.1|   0.0|  0.12|
// |  0|  1|  1| 0.3|    2|   0.0|  0.11|  0.12|
// +---+---+---+----+-----+------+------+------+