Spark：规范化数据帧的每一行

Question

我有一个 Spark DataFrame，如下所示

df.show()
+------+------+------+
|  col1|  col2|  col3|
+------+------+------+
|   5.0|   5.0|   0.0|
|   2.0|   3.0|   5.0|
|   4.0|   1.0|  10.0|
+------+------+------+

我想标准化每个单独的行，以便在操作后，新列看起来像：

+--------+--------+--------+
|new_col1|new_col2|new_col3|
+--------+--------+--------+
|     0.5|     0.5|     0.0|
|     0.2|     0.3|     0.5|
|0.266667|0.066667|0.666667|
+--------+--------+--------+

更正式地说，我想申请的操作是：

对于每一行，

    new_col_i = col_i / (col_1 + col_2 + col_3)

我需要以编程方式执行此操作，而不是列出所有列，因为我的 DataFrame 有很多列。

当前解决方案：

我想到的当前解决方案是创建一个列来表示每行所有条目的总和，然后将每列除以该总和列。

var newDF = df.withColumn("total", df.columns.map(c => col(c)).reduce((c1, c2) => c1 + c2))

for (c <- Array("col1", "col2", "col3")) {
    newDF = newDF.withColumn("normalized_" + c, col(c).divide(col("total")))
}
newDF.show()

+----+----+----+-----+-------------------+-------------------+------------------+
|col1|col2|col3|total|    normalized_col1|    normalized_col2|   normalized_col3|
+----+----+----+-----+-------------------+-------------------+------------------+
| 5.0| 5.0| 0.0| 10.0|                0.5|                0.5|               0.0|
| 2.0| 3.0| 5.0| 10.0|                0.2|                0.3|               0.5|
| 4.0| 1.0|10.0| 15.0|0.26666666666666666|0.06666666666666667|0.6666666666666666|
+----+----+----+-----+-------------------+-------------------+------------------+

任何使代码更简洁的替代方法？

Answer 1

您的解决方案是正确的，不能改进太多。您可以var通过将 for 循环替换为 a来摆脱不惯用的用法foldLeft，并在周围使用更多的语法糖，但除此之外它会保持不变：

val withTotal = df.withColumn("total", df.columns.map(col).reduce(_ + _))

val result = df.columns.foldLeft(withTotal) {
  (tmp, c) => tmp.withColumn(s"new_$c", $"$c" / $"total")
}
  .drop(df.columns: _*)
  .drop("total")