Spark scala 從列表中選擇多列和單列

Question

我正在嘗試在 dataframe 上執行 select 但我遇到了一些麻煩。

我有這個初始 dataframe

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

而我要做的是將 value_a 與 value_b 相加並保持其他相同。 所以我有這個清單

val select_list = List(id, value_c, value_d)

然后我做 select

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)
+----------+-------+-------+

但我得到“這里允許沒有 _* 注釋”。 請記住，實際上我有很多列，所以我需要使用一個列表，我不能簡單地 select 每一列。 我遇到了這個麻煩，因為作為 sum 的結果的新列與現有列具有相同的名稱，所以我不能只 select(column("*"), sum....).drop (value_b) 或者我將刪除舊列和帶有總和的新列。

在單個 select 中添加多個和單個列的正確語法是什么，或者我還能如何解決這個問題？ 現在我決定這樣做：

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 很昂貴，因為我正在創建一個新的 dataframe 並帶有一個新的或重命名的列，並且我正在尋找成本更低的操作。

提前致謝！

Answer 1

只需使用.withColumn function，它將替換列（如果存在）：

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|
+------+------+------+---+