sparkscala从列表中选择多列和单列

ncecgwcz  于 2021-05-16  发布在  Spark
关注(0)|答案(2)|浏览(375)

我正在尝试对Dataframe进行选择,但遇到了一点问题。
我有这个初始Dataframe

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

我要做的是把值a和值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),否则我将删除旧列和包含该和的新列。
在一个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,其中包含一个新的或重命名的列,我正在寻找成本更低的操作。
提前谢谢!

fkaflof6

fkaflof61#

简单使用 .withColumn 函数,它将替换列(如果存在):

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

3zwjbxry2#

您可以创建一个新的sum字段,并收集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|
+------+------+------+---+

相关问题