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中的每个列，foe每个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

我所做的是通过连接数据集 2 映射数据集 1 中的每一行

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: _*)