In PySpark, adding a new column to a DataFrame is a common and essential operation, often used for transforming data, performing calculations, or enriching the dataset. PySpark offers 3 main methods for this: withColumn(),select() and selectExpr(). These methods allow you to create new columns, but they serve different purposes and are used in different contexts.

This article will guide you through adding new columns using these 3 methods, explaining their use cases and providing examples.

Method 1: Using withColumn(): This method is the most common way to add or update a column in a PySpark DataFrame. It is straightforward and allows you to create a new column by applying an expression to existing columns.

Syntax: DataFrame.withColumn(colName, col)
arguments:
colName: The name of the new or existing column.
col: A PySpark expression, which can be an operation on existing columns or a constant value.

Example 1: Add a New Column by Applying an Operation: Let’s start by creating a simple DataFrame and then adding a new column based on a calculation.

# Importing necessary modules
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# Initialize Spark session
spark = SparkSession.builder.appName("Add Column Example").getOrCreate()

# Sample DataFrame
data = [("Alice", 25), ("Bob", 30), ("Catherine", 29)]
columns = ["Name", "Age"]
df = spark.createDataFrame(data, columns)

# Add a new column 'Age_After_5_Years' using withColumn
df_with_new_column = df.withColumn("Age_After_5_Years", col("Age") + 5)

# Show the result
df_with_new_column.show()

Output:

+----------+---+-------------------+
|      Name|Age| Age_After_5_Years |
+----------+---+-------------------+
|     Alice| 25|                 30|
|       Bob| 30|                 35|
| Catherine| 29|                 34|
+----------+---+-------------------+

Explanation:

• withColumn(): Used to add a new column.
• col(“Age”) + 5: Adds 5 to each value in the “Age” column and stores the result in the new column “Age_After_5_Years”.

Example 2: Add a New Column with a Constant Value: You can also use withColumn() to add a column with a constant or literal value across all rows.

from pyspark.sql.functions import lit

# Add a constant value column
df_with_constant = df.withColumn("Country", lit("INDIA"))

# Show the result
df_with_constant.show()

Output:

+----------+---+-------+
|      Name|Age|Country|
+----------+---+-------+
|     Alice| 25|  INDIA|
|       Bob| 30|  INDIA|
| Catherine| 29|  INDIA|
+----------+---+-------+

Explanation:

• lit(): This function creates a constant literal value (“INDIA”) that is assigned to the new column “Country.”

Method 2: Using select() to Add a New Column:While withColumn() is designed to add a single new column, the select() method can also be used to create new columns. It’s especially useful when you need to select specific columns and apply transformations at the same time.

Syntax:DataFrame.select(*cols)
arguments:
cols: A list of columns or expressions to select. You can select existing columns or apply transformations to create new ones.

Example 1: Add a New Column While Selecting Specific Columns: In this example, we’ll use the select() method to choose certain columns from the DataFrame and create a new column as part of the process.

# Add a new column 'Age_After_5_Years' using select
df_with_select = df.select("Name", "Age", (col("Age") + 5).alias("Age_After_5_Years"))

# Show the result
df_with_select.show()

Output:

+----------+---+-------------------+
|      Name|Age|Age_After_5_Years  |
+----------+---+-------------------+
|     Alice| 25|                 30|
|       Bob| 30|                 35|
| Catherine| 29|                 34|
+----------+---+-------------------+

Explanation:

• select(): Used to select columns and apply transformations at the same time.
• alias(): Renames the column created by the expression col(“Age”) + 5 to “Age_After_5_Years”.

Example 2: Add Multiple New Columns Using select(): If you need to add multiple new columns at once, select() can be more concise than calling withColumn() multiple times.

# Add multiple new columns using select
df_with_multiple = df.select(
    "Name",
    "Age",
    (col("Age") + 5).alias("Age_After_5_Years"),
    (col("Age") * 365).alias("Age_in_Days")
)

# Show the result
df_with_multiple.show()

Output:

+----------+---+-------------------+-----------+
|      Name|Age|Age_After_5_Years  |Age_in_Days|
+----------+---+-------------------+-----------+
|     Alice| 25|                30|        9125|
|       Bob| 30|                35|       10950|
| Catherine| 29|                34|       10585|
+----------+---+-------------------+-----------+

Explanation: This example shows how select() can be used to add multiple new columns with different transformations in one step.

• The first new column, “Age_After_5_Years” adds 5 to the age.
• The second new column, “Age_in_Days” multiplies the age by 365 to convert it to the number of days.

Method 3: Using selectExpr() to Add a New Column: This method allows you to specify SQL-like expressions, which means you can use it to add new columns by defining them within the expression.

Syntax: DataFrame.selectExpr(*exprs)
arguments:
exprs: One or more SQL-like expressions. You can include both existing columns and new ones derived from expressions.

Example 1: Adding a New Column : Let’s say you have a DataFrame containing columns Name and Age, and you want to add a new column called AgeInMonths, which is calculated as Age * 12.

# Add a new column 'AgeInMonths' using selectExpr
df_with_selectExpr= df.selectExpr("Name", "Age", "Age * 12 as AgeInMonths")

# Show the result
df_with_selectExpr.show()

Output:

+----------+---+-----------+
|      Name|Age|AgeInMonths|
+----------+---+-----------+
|     Alice| 25|        300|
|       Bob| 30|        360|
| Catherine| 29|        348|
+----------+---+-----------+

Explanation:

• The original columns Name and Age are retained.The new column AgeInMonths is created by multiplying the Age column by 12.

Comparing withColumn(), select() and selectExpr():

These methods are powerful tools for adding new columns, but they have different strengths:

• withColumn(): This method is ideal for adding or updating one column at a time. It’s simple to use and can be called repeatedly to modify or create additional columns.
• select(): This method is useful when you want to perform multiple transformations at once or when you need to select specific columns while adding new ones.
• selectExpr(): It is more powerful when you want to add multiple columns or when you prefer SQL-like syntax. It’s also more flexible when performing transformations or conditional logic.

Conclusion: Adding new columns to a DataFrame is a crucial part of transforming and analyzing data in PySpark. In this article, we explored three primary methods for adding columns. Understanding when and how to use these methods will help you efficiently transform your DataFrames in PySpark. Whether you’re performing simple calculations, creating derived columns, or applying complex transformations, PySpark provides the tools to handle your data processing needs.