Introduction

PySpark withColumn() is a DataFrame function used to add a new column to the DataFrame, change the value of an existing column, convert the data type of a column, or derive a new column from an existing column.

PySpark withColumn() Syntax and Usage

PySpark withColumn() is a transformation function of a DataFrame used to manipulate the column values of all rows or selected rows in the DataFrame.

The withColumn() function returns a new PySpark DataFrame after performing operations like adding a new column, updating the value of an existing column, deriving a new column from an existing column, and more.

Below is the syntax of withColumn() function.

withColumn(colName, col)

colName – specify the new column you want to create. Use an existing column to update the value.

col – column expression.

Since withColumn() is a transformation function, it doesn’t execute until an action is called.

PySpark withColumn() method introduces a projection internally. Therefore, calling it multiple times, for instance, via loops in order to add multiple columns can generate complex plans which can cause performance issues. To avoid this, use select with multiple columns at once as suggested in the Spark Documentation.

First, let’s create a DataFrame to work with.

from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.sql.types import StructType, StructField, StringType

spark = SparkSession.builder.appName('ExampleApp').getOrCreate()

data = [
  Row(name=Row(firstname="James", middlename="", lastname="Smith"), dob="36636", gender="M", salary="3000"),
  Row(name=Row(firstname="Michael", middlename="Rose", lastname=""), dob="40288", gender="M", salary="4000"),
  Row(name=Row(firstname="Robert", middlename="", lastname="Williams"), dob="42114", gender="M", salary="4000"),
  Row(name=Row(firstname="Maria", middlename="Anne", lastname="Jones"), dob="39192", gender="F", salary="4000"),
  Row(name=Row(firstname="Jen", middlename="Mary", lastname="Brown"), dob="", gender="F", salary="-1")
]

schema = StructType([
    StructField('name', StructType([
        StructField('firstname', StringType(), True),
        StructField('middlename', StringType(), True),
        StructField('lastname', StringType(), True)
    ])),
    StructField('dob', StringType(), True),
    StructField('gender', StringType(), True),
    StructField('salary', StringType(), True)
])

df = spark.createDataFrame(data, schema)

1. Add a New Column to DataFrame

To create a new column, pass your desired column name to the first argument of withColumn() transformation function. Ensure this new column is not already present in DataFrame; if it is, it updates the value of the column. In the snippet below, the lit() function is used to add a constant value to a DataFrame column. We can also chain operations to add multiple columns.

from pyspark.sql.functions import lit

# Add a New Column to DataFrame
df.withColumn("Country", lit("USA")).show()

# Chaining to operate on multiple columns
df.withColumn("Country", lit("USA")) \
  .withColumn("anotherColumn", lit("anotherValue")).show()

This approach is fine if you are manipulating a few columns, but when you want to add or update multiple columns, do not chain withColumn() as it leads to performance issues. Instead, use select() to update multiple columns.

2. Change Value of an Existing Column

withColumn() function can also be used to update the value of an existing column. To change the value, pass an existing column name as the first argument and a value to be assigned as a second column. Note that the second argument should be a Column type.

from pyspark.sql.functions import col

# Change Value of an Existing Column
df.withColumn("salary", col("salary") * 100).show()

This snippet multiplies the value of “salary” by 100 and updates it back to the “salary” column.

3. Derive a New Column From an Existing Column

To create a new column, specify the first argument with a name you want your new column to be and use the second argument to assign a value by applying an operation on an existing column.

# Derive New Column From an Existing Column
df.withColumn("CopiedColumn", col("salary") * -1).show()

This snippet creates a new column, “CopiedColumn,” by multiplying the “salary” column with value -1.

4. Change Column Data Type

By using PySpark withColumn on a DataFrame and using the cast function on a column, we can change the datatype of a DataFrame column. The below statement changes the datatype from String to Integer for the “salary” column.

# Change Column Data Type
df.withColumn("salary", col("salary").cast("Integer")).show()

5. Split Column into Multiple Columns

Though this example doesn’t use withColumn() function, it’s useful to explain splitting one DataFrame column into multiple columns using PySpark transformations.

from pyspark.sql import functions as F

# Split Column into Multiple Columns
columns = ["name", "address"]
data2 = [
  ("Robert, Smith", "1 Main st, Newark, NJ, 92537"),
  ("Maria, Garcia", "3456 Walnut st, Newark, NJ, 94732")
]

dfFromData = spark.createDataFrame(data2, columns)
dfFromData.select(F.split(F.col("name"), ",").alias("Name"), 
                  F.split(F.col("address"), ",").alias("Address")) \
          .select(
              F.col("Name").getItem(0).alias("First Name"),
              F.col("Name").getItem(1).alias("Last Name"),
              F.col("Address").getItem(0).alias("Address Line1"),
              F.col("Address").getItem(1).alias("City"),
              F.col("Address").getItem(2).alias("State"),
              F.col("Address").getItem(3).alias("ZipCode")
          ).show(truncate=False)

This snippet splits the name column into “First Name” and “Last Name,” and the address column into “Address Line1,” “City,” “State,” and “ZipCode.”

The output of the final DataFrame will be:

+----------+---------+--------------+-------+-----+-------+
|First Name|Last Name|Address Line1 |City   |State|ZipCode|
+----------+---------+--------------+-------+-----+-------+
|Robert    |    Smith|1 Main st     | Newark|  NJ |  92537|
|Maria     |   Garcia|3456 Walnut st| Newark|  NJ |  94732|
+----------+---------+--------------+-------+-----+-------+