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

Posted 2023-04-15

【中文标题】Spark scala 从列表中选择多列和单列

【英文标题】：Spark scala selecting multiple columns from a list and single columns

【发布时间】：2020-11-27 15:55:20

【问题描述】：

我正在尝试对数据框进行选择，但遇到了一些麻烦。

我有这个初始数据框

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

但我得到“这里允许没有 _* 注释”。请记住，实际上我有很多列，所以我需要使用一个列表，我不能简单地选择每一列。我遇到了这个麻烦，因为作为 sum 的结果的新列与现有列具有相同的名称，所以我不能只 select(column("*"), sum....).drop (value_b) 否则我会删除旧列和带有总和的新列。

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

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 很昂贵，因为我正在创建一个带有新列或重命名列的新数据框，并且我正在寻找成本更低的操作。

提前致谢！

【问题讨论】：

仅供参考，您的 (select_list.map(col):_* 被禁止的原因是 _* 注释必须在末尾。如果您将总和列放在列表之前，它会起作用。

【参考方案1】：

只需使用.withColumn函数，如果存在则替换列：

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

【参考方案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|
+------+------+------+---+