scala - Spark scala 中数据帧行内容的条件

Question

我有以下数据框:

+--------+---------+------+
|  value1| value2  |value3|
+--------+---------+------+
|   a    |  2      |   3  |
+--------+---------+------+
|   b    |  5      |   4  |
+--------+---------+------+
|   b    |  5      |   4  |
+--------+---------+------+
|   c    |  3      |   4  |
+--------+---------+------+

当value1=b时，我想把行的value2/value3的结果放在中。所有行 (甚至不属于 b 的行)在名称为“结果”的新字段中。这意味着必须将另一列添加到数据框中。例如，对于所有行，应将 5/4 的结果(我选择它，因为它用于 b)应添加到数据帧中。我知道，我应该使用以下代码:

 val dataframe_new = Dataframe.withColumn("result", $"value1" / $"value2")
 Dataframe.show()

但是，我如何才能将条件添加到所有行中。输出应如下所示:

+---+---+---+------+
| v1| v2| v3|result|
+---+---+---+------+
|  a|  2|  3|  1.25|
|  b|  5|  4|  1.25|
|  b|  5|  4|  1.25|
|  c|  3|  4|  1.25|
+---+---+---+------+

你能帮助我吗？提前致谢。

Answer 1

您只需要使用 when :

scala> val df = Seq(("a",2,3),("b",5,4),("b",5,4),("c",3,4)).toDF("v1","v2","v3")
df: org.apache.spark.sql.DataFrame = [v1: string, v2: int ... 1 more field]

scala> df.withColumn("result", when($"v1" === "b" , ($"v2"/$"v3"))).show
+---+---+---+------+
| v1| v2| v3|result|
+---+---+---+------+
|  a|  2|  3|  null|
|  b|  5|  4|  1.25|
|  b|  5|  4|  1.25|
|  c|  3|  4|  null|
+---+---+---+------+

您可以嵌入多个 when如下:

scala> df.withColumn("result", when($"v1" === "b" , ($"v2"/$"v3")).
     |    otherwise(when($"v1" === "a", $"v3"/$"v2"))).show
+---+---+---+------+
| v1| v2| v3|result|
+---+---+---+------+
|  a|  2|  3|   1.5|
|  b|  5|  4|  1.25|
|  b|  5|  4|  1.25|
|  c|  3|  4|  null|
+---+---+---+------+

编辑:似乎您需要其他的东西来满足 v1 的条件。始终具有相同的值 v2和 v3这允许我们执行以下操作:

与 Spark 2+ :

scala> val res = df.filter($"v1" === lit("b")).distinct.select($"v2"/$"v3").as[Double].head
res: Double = 1.25

之前 Spark <2 :

scala> val res = df.filter($"v1" === lit("b")).distinct.withColumn("result",$"v2"/$"v3").rdd.map(_.getAs[Double]("result")).collect()(0)
res: Double = 1.25                                                              

scala> df.withColumn("v4", lit(res)).show
+---+---+---+----+
| v1| v2| v3|  v4|
+---+---+---+----+
|  a|  2|  3|1.25|
|  b|  5|  4|1.25|
|  b|  5|  4|1.25|
|  c|  3|  4|1.25|
+---+---+---+----+