How to Add a New Column to a Spark DataFrame Using PySpark?

Adding a new column to a Spark DataFrame in PySpark is a common operation you might need in data processing. You can achieve this in several ways, depending on your specific needs. Below, I’ll explain a couple of methods, along with code snippets and their expected output.

Method 1: Using the `withColumn` Method

The `withColumn` method is one of the most straightforward ways to add a new column to a DataFrame. You can create a new column based on existing columns or provide a constant value.

Step-by-Step Explanation

1. Import the necessary libraries.
2. Create a sample DataFrame.
3. Add a new column using the `withColumn` method.

Code Snippet

from pyspark.sql import SparkSession
from pyspark.sql.functions import lit

# Create a Spark session
spark = SparkSession.builder.appName("AddColumnExample").getOrCreate()

# Define a sample DataFrame
data = [("Alice", 34), ("Bob", 45), ("Catherine", 29)]
columns = ["Name", "Age"]
df = spark.createDataFrame(data, columns)

# Add a new column "Country" with a constant value "USA"
df_with_new_column = df.withColumn("Country", lit("USA"))

# Show the resulting DataFrame
df_with_new_column.show()

Output

+---------+---+-------+
|     Name|Age|Country|
+---------+---+-------+
|    Alice| 34|    USA|
|      Bob| 45|    USA|
| Catherine| 29|    USA|
+---------+---+-------+

Method 2: Using the `select` Method

If you need to add a new column based on some computation of existing columns, you can use the `select` method in combination with other PySpark SQL functions.

Step-by-Step Explanation

1. Import the necessary libraries.
2. Create a sample DataFrame.
3. Use the `select` method to compute the new column.

Code Snippet

from pyspark.sql import functions as F

# Create a sample DataFrame
data = [("Alice", 34), ("Bob", 45), ("Catherine", 29)]
columns = ["Name", "Age"]
df = spark.createDataFrame(data, columns)

# Add a new column "Age_in_5_years" by adding 5 to the "Age" column
df_with_new_column = df.select("Name", "Age", (F.col("Age") + 5).alias("Age_in_5_years"))

# Show the resulting DataFrame
df_with_new_column.show()

Output

+---------+---+-------------+
|     Name|Age|Age_in_5_years|
+---------+---+-------------+
|    Alice| 34|           39|
|      Bob| 45|           50|
| Catherine| 29|           34|
+---------+---+-------------+

Conclusion

These are two common methods to add a new column to a Spark DataFrame using PySpark. The `withColumn` method is more straightforward for adding columns with constant values or simple expressions, while the `select` method is useful for more complex computations. Choose the method that best suits your use case.