Spark scala从列表中选择多个列和单个列

Question

我试图对数据帧进行选择，但遇到了一点麻烦。

我有这个初始数据帧

+----------+-------+-------+-------+
|id|value_a|value_b|value_c|value_d|
+----------+-------+-------+-------+

我要做的是将value_a和value_b相加，并保持其他值不变。所以我有这个列表

val select_list = List(id, value_c, value_d)

然后我选择

df.select(select_list.map(col):_*, (col(value_a) + col(value_b)).as("value_b"))

我期待得到这个：

+----------+-------+-------+
|id|value_c|value_d|value_b|  --- that value_b is the sum of value_a and value_b (original)
+----------+-------+-------+

但我得到了"；此处不允许注释"；。请记住，实际上我有很多列，所以我需要使用列表，我不能简单地选择每一列。我遇到了这个麻烦，因为作为求和结果的新列与现有列的名称相同，所以我不能只选择(column("*"(，sum…(。drop(value_b(，否则我将删除旧列和带有sum的新列。

在一次选择中添加多列和单列的正确语法是什么？或者我如何解决这个问题？现在我决定这样做：

df.select(col("*"), (col(value_a) + col(value_b)).as("value_b_tmp")).
drop("value_a", "value_b").withColumnRenamed("value_b_tmp", "value_b")

这很好，但我知道withColumn和withColumnRenamed的成本很高，因为我正在用一个新的或重命名的列创建一个几乎新的数据帧，我正在寻找成本更低的操作。

提前感谢！

Answer 1

只需使用.withColumn函数，它就会替换列(如果存在(：

df
.withColumn("value_b", col("value_a") + col("value_b"))
.select(select_list.map(col):_*)

Answer 2

您可以创建一个新的求和字段，并将n列之和的运算结果收集为：

val df: DataFrame = 
spark.createDataFrame(
spark.sparkContext.parallelize(Seq(Row(1,2,3),Row(1,2,3))),
StructType(List(
StructField("field1", IntegerType), 
StructField("field2", IntegerType), 
StructField("field3", IntegerType))))
val columnsToSum = df.schema.fieldNames
columnsToSum.filter(name =>  name != "field1")
.foldLeft(df.withColumn("sum", lit(0)))((df, column) =>
df.withColumn("sum", col("sum") + col(column)))

提供：

+------+------+------+---+
|field1|field2|field3|sum|
+------+------+------+---+
|     1|     2|     3|  5|
|     1|     2|     3|  5|
+------+------+------+---+