Search code examples
apache-sparkpysparkapache-spark-sqlrename

Rename more than one column using withColumnRenamed


I want to change names of two columns using spark withColumnRenamed function. Of course, I can write:

data = sqlContext.createDataFrame([(1,2), (3,4)], ['x1', 'x2'])
data = (data
       .withColumnRenamed('x1','x3')
       .withColumnRenamed('x2', 'x4'))

but I want to do this in one step (having list/tuple of new names). Unfortunately, neither this:

data = data.withColumnRenamed(['x1', 'x2'], ['x3', 'x4'])

nor this:

data = data.withColumnRenamed(('x1', 'x2'), ('x3', 'x4'))

is working. Is it possible to do this that way?


Solution

  • It is not possible to use a single withColumnRenamed call.

    • You can use DataFrame.toDF method*

      data.toDF('x3', 'x4')
      

      or

      new_names = ['x3', 'x4']
      data.toDF(*new_names)
      
    • It is also possible to rename with simple select:

      from pyspark.sql.functions import col
      
      mapping = dict(zip(['x1', 'x2'], ['x3', 'x4']))
      data.select([col(c).alias(mapping.get(c, c)) for c in data.columns])
      

    Similarly in Scala you can:

    • Rename all columns:

      val newNames = Seq("x3", "x4")
      
      data.toDF(newNames: _*)
      
    • Rename from mapping with select:

      val  mapping = Map("x1" -> "x3", "x2" -> "x4")
      
      df.select(
        df.columns.map(c => df(c).alias(mapping.get(c).getOrElse(c))): _*
      )
      

      or foldLeft + withColumnRenamed

      mapping.foldLeft(data){
        case (data, (oldName, newName)) => data.withColumnRenamed(oldName, newName) 
      }
      

    * Not to be confused with RDD.toDF which is not a variadic functions, and takes column names as a list,