Home > Software design >  How to rename column of a DataFrame from a list
How to rename column of a DataFrame from a list

Time:03-25

I have a dataframe like this:

dog cat
Cell 1 Cell 2
Cell 3 Cell 4

And a list like this:

dog, bulldog

cat, persian

I would like to create a function that find the name of the column in the list and substitute it with the second element (bulldog, persian).

So the final result should be:

| bulldog  | persian  |
| -------- | -------- |
| Cell 1   | Cell 2   |
| Cell 3   | Cell 4   |

CodePudding user response:

You need to perform a look-up for your original column in the pre-defined list that you have shown. It's easier to create a Map out of it so lookups can be performed:

val list: List[(String, String)] = List(("dog", "bulldog"), ("cat", "persian"))

val columnMap = list.toMap

// columnMap: scala.collection.immutable.Map[String,String] = Map(dog -> bulldog, cat -> persian)


val originalCols = df.columns
​
val renamedCols = originalCols.map{
  c => if (columnMap.keys.toArray.contains(c)) s"${c} as ${columnMap.getOrElse(c, "")}"
       else c
}

println(renamedCols)

// renamedCols: Array[String] = Array(dog as bulldog, cat as persian)

df.selectExpr(renamedCols: _*).show(false)

//  ------- ------- 
// |bulldog|persian|
//  ------- ------- 
// |Cell 1 |Cell 2 |
// |Cell 1 |Cell 2 |
//  ------- ------- 
  • Related