How To Add a New Column To a PySpark DataFrame | Towards Data Science

Adding new columns to PySpark DataFrames is probably one of the most common operations you need to perform as part of your day-to-day work.
In today’s short guide, we will discuss about how to do so in many different ways. Specifically, we will explore how to add new columns and populate them
with literals
by transforming existing columns
using joins
using functions or UDFs
First, let’s create an example DataFrame that we’ll reference throughout this article to demonstrate the concepts we are interested in
from pyspark.sql import SparkSession# Create an instance of spark sessionspark_session = SparkSession.builder \    .master('local[1]') \    .appName('Example') \    .getOrCreate()# Create an example DataFramedf = spark_session.createDataFrame(    [        (1, True, 'a', 1.0),        (2, True, 'b', 2.0),        (3, False, None, 3.0),        (4, False, 'd', None),    ],    ['colA', 'colB', 'colC', 'colD'])df.show()+----+-----+----+----+|colA| colB|colC|colD|+----+-----+----+----+|   1| true|   a| 1.0||   2| true|   b| 2.0||   3|false|null| 3.0||   4|false|   d|null|+----+-----+----+----+
Assuming that you want to add a new column containing literals, you can make use of the pyspark.sql.functions.lit function that is used to create a column of literals.
For example, the following command will add a new column called colE containing the value of 100 in each row.
df.withColumn('colE', lit(100))df.show()+----+-----+----+----+----+|colA| colB|colC|colD|colE|+----+-----+----+----+----+|   1| true|   a| 1.0| 100||   2| true|   b| 2.0| 100||   3|false|null| 3.0| 100||   4|false|   d|null| 100|+----+-----+----+----+----+
Note that you have to use lit function because the second argument of withColumn must be of type Column.
Now if you want to add a column containing more complex data structures such as an array, you can do so as shown below:
from pyspark.sql.functions import lit, arraydf = df.withColumn('colE', array(lit(100), lit(200), lit(300)))df.show()+----+-----+----+----+---------------+|colA| colB|colC|colD|           colE|+----+-----+----+----+---------------+|   1| true|   a| 1.0|[100, 200, 300]||   2| true|   b| 2.0|[100, 200, 300]||   3|false|null| 3.0|[100, 200, 300]||   4|false|   d|null|[100, 200, 300]|+----+-----+----+----+---------------+
If you want to create a new column based on an existing column then again you should specify the desired operation in withColumn method.
For example, if you want to create a new column by multiplying the values of an existing column (say colD) with a constant (say 2), then the following will do the trick:
from pyspark.sql.functions import coldf = df.withColumn('colE', col('colD') * 2)df.show()+----+-----+----+----+----+|colA| colB|colC|colD|colE|+----+-----+----+----+----+|   1| true|   a| 1.0| 2.0||   2| true|   b| 2.0| 4.0||   3|false|null| 3.0| 6.0||   4|false|   d|null|null|+----+-----+----+----+----+
Alternatively, we can still create a new DataFrame and join it back to the original one. First, you need to create a new DataFrame containing the new column you want to add along with the key that you want to join on the two DataFrames
new_col = spark_session.createDataFrame(    [(1, 'hello'), (2, 'hi'), (3, 'hey'), (4, 'howdy')],    ('key', 'colE'))new_col.show()+---+-----+|key| colE|+---+-----+|  1|hello||  2|   hi||  3|  hey||  4|howdy|+---+-----+
and finally do the join:
from pyspark.sql.functions import coldf = df \    .join(new_col, col('colA') == col('key'), 'leftouter') \    .drop('key')df.show()+----+-----+----+----+-----+|colA| colB|colC|colD| colE|+----+-----+----+----+-----+|   1| true|   a| 1.0|hello||   2| true|   b| 2.0|   hi||   3|false|null| 3.0|  hey||   4|false|   d|null|howdy|+----+-----+----+----+-----+
Another possibility, is to use a function that returns a Column and pass that function to withColumn. For instance, you can use the built-in pyspark.sql.functions.rand function to create a column containing random numbers, as shown below:
from pyspark.sql.functions import randdf = df.withColumn('colE', rand())df.show()+----+-----+----+----+--------------------+|colA| colB|colC|colD|                colE|+----+-----+----+----+--------------------+|   1| true|   a| 1.0|0.026110187082684866||   2| true|   b| 2.0|0.046264104329627576||   3|false|null| 3.0|  0.7892572670252188||   4|false|   d|null|  0.7963792998818318|+----+-----+----+----+--------------------+
In today’s short guide we discussed about how to insert additional columns to existing PySpark DataFrames.
Become a member and read every story on Medium. Your membership fee directly supports me and other writers you read. You’ll also get full access to every story on Medium.
You may also like