Programmatically Rename All But One Column Spark Scala

Viewed 1150

I have this DataFrame:

val df = Seq(
   ("LeBron", 36, 18, 12),
   ("Kevin", 42, 8, 9),
   ("Russell", 44, 5, 14)).
   toDF("player", "points", "rebounds", "assists")

df.show()
+-------+------+--------+-------+
| player|points|rebounds|assists|
+-------+------+--------+-------+
| LeBron|    36|      18|     12|
|  Kevin|    42|       8|      9|
|Russell|    44|       5|     14|
+-------+------+--------+-------+

I want to add "season_high" to every column name except player. I also want to use a function to do this because my real data set has 250 columns.

I've come up with the method below that gets me the output that I want, but I'm wondering if there is a way to pass a rule to the renamedColumns mapping function that makes it so that the column name player doesn't get switched to season_high_player, and then back to player with the additional .withColumnRenamed function.

val renamedColumns = df.columns.map(name => col(name).as(s"season_high_$name"))

val df2 = df.select(renamedColumns : _*).
    withColumnRenamed("season_high_player", "player")

df2.show()
+-------+------------------+--------------------+-------------------+
| player|season_high_points|season_high_rebounds|season_high_assists|
+-------+------------------+--------------------+-------------------+
| LeBron|                36|                  18|                 12|
|  Kevin|                42|                   8|                  9|
|Russell|                44|                   5|                 14|
+-------+------------------+--------------------+-------------------+
3 Answers

@philantrovert was right but he just forgot to tell you how to use that "formula", so here you go :

val selection : Seq[Column] = Seq(col("player")) ++ df.columns.filter(_ != "player")
                                        .map(name => col(name).as(s"season_high_$name"))
df.select(selection : _*).show
// +-------+------------------+--------------------+-------------------+
// | player|season_high_points|season_high_rebounds|season_high_assists|
// +-------+------------------+--------------------+-------------------+
// | LeBron|                36|                  18|                 12|
// |  Kevin|                42|                   8|                  9|
// |Russell|                44|                   5|                 14|
// +-------+------------------+--------------------+-------------------+

So what we have done here is to filter out the column name that we don't need (This is plain scala). Then we map the column names that we kept to convert them into column which we rename.

you can do the following by making the one column you don't want to rename as the first column and apply the following logic

import org.apache.spark.sql.functions._
val columnsRenamed = col(df.columns.head) +: df.columns.tail.map(name => col(name).as(s"season_high_$name"))
df.select(columnsRenamed :_*).show(false)

You should be getting output as

+-------+------------------+--------------------+-------------------+
|player |season_high_points|season_high_rebounds|season_high_assists|
+-------+------------------+--------------------+-------------------+
|LeBron |36                |18                  |12                 |
|Kevin  |42                |8                   |9                  |
|Russell|44                |5                   |14                 |
+-------+------------------+--------------------+-------------------+

one more variation that doesn't depend on the position of the field.

scala>     val df = Seq(
     |       ("LeBron", 36, 18, 12),
     |       ("Kevin", 42, 8, 9),
     |       ("Russell", 44, 5, 14)).
     |       toDF("player", "points", "rebounds", "assists")
df: org.apache.spark.sql.DataFrame = [player: string, points: int ... 2 more fields]

scala> val newColumns = df.columns.map( x => x match { case "player" => col("player") case x =>  col(x).as(s"season_high_$x")}  )
newColumns: Array[org.apache.spark.sql.Column] = Array(player, points AS `season_high_points`, rebounds AS `season_high_rebounds`, assists AS `season_high_assists`)

scala> df.select(newColumns:_*).show(false)
+-------+------------------+--------------------+-------------------+
|player |season_high_points|season_high_rebounds|season_high_assists|
+-------+------------------+--------------------+-------------------+
|LeBron |36                |18                  |12                 |
|Kevin  |42                |8                   |9                  |
|Russell|44                |5                   |14                 |
+-------+------------------+--------------------+-------------------+


scala>
Related