在 Spark 2.0 中进行列级操作的有效方法

Question

我有两个数据集

数据集 1：

id  a    b      c     d
1  0.3  0.1   0.2   0.2
2  0.2  0.3   0.3   0.4
3  0.2  0.4   0.7   0.7
....

数据集2

id  x    
1   8  
2   4 
3   10  
....

我想做一个操作，使用dataset2中的“x”列乘以dataset1中的每一列，每个id，这样期望的输出是：

id   a    b    c    d
1   2.4   0.8  1.6  1.6
2   0.8   1.2  1.2  1.6
3    2     4    7    7

我所做的是通过加入 dataset2 来映射 dataset1 中的每一行

val result = dataset1.join(dataset2, Seq("id")
                     .map(row=> row.getAs[String]("id"),
                          row=> row.getAs[Double]("a") * row.getAs[Int]("x"),
                          row=> row.getAs[Double]("b") * row.getAs[Int]("x"),
                          row=> row.getAs[Double]("c") * row.getAs[Int]("x"),
                          row=> row.getAs[Double]("d") * row.getAs[Int]("x"))

我觉得这样写有点多余。有什么办法可以更清楚吗？

Answer 1

您只需要select:

dataset1.join(dataset2, Seq("id")).select(
  $"id", $"a" * $"x", $"b" * $"x", $"c" * $"x", $"d" * "x"
).toDF("id", "a", "b", "c", "d")

可以泛化

val exprs = $"id" +: dataset1.columns.tail.map(c => (col(c) * $"x").alias(c))
dataset1.join(dataset2, Seq("id")).select(exprs: _*)