PySpark: Dataframe Add Columns

Pradeep

PySpark: Dataframe Add Columns

This tutorial will explain various approaches with examples on how to add new columns or modify existing columns in a dataframe. Below listed topics will be explained with examples on this page, click on item in the below list and it will take you to the respective section of the page:

List all Columns
Add Column using withColumn
Add Column using select

Add Column using other dataframe

Sample Data: Dataset used in the below examples can be downloaded from here .


df = spark.read.csv("file:///path_to_files/csv_file_with_duplicates.csv", header=True)

df.show()
+-----+---------+-------+
|db_id|  db_name|db_type|
+-----+---------+-------+
|   12| Teradata|  RDBMS|
|   14|Snowflake|CloudDB|
|   15|  Vertica|  RDBMS|
|   12| Teradata|  RDBMS|
|   22|    Mysql|  RDBMS|
+-----+---------+-------+

➠ List all Columns: columns attribute can be used on a dataframe to return all the column names as a list.


df.columns

Output: ['db_id', 'db_name', 'db_type']

➠ Add Column using withColumn: withColumn() function can be used on a dataframe to either add a new column or replace an existing column that has same name. withColumn() function can cause performance issues and even "StackOverflowException" if it is called multiple times using loop to add multiple columns. Spark suggests to use "select" function to add multiple columns at once.

Syntax:
```
withColumn(columnName, columnLogic/columnExpression)
```
This function takes 2 parameters, 1st parameter is the name of new or existing column and 2nd parameter is the column logic / column expression for the new or existing column passed in the 1st paramter.

Example 1: Value of existing column db_type is updated to literal string "Relation Database" for all the rows.


from pyspark.sql.functions import col,lit
df_updated = df.withColumn("db_type",lit("Relation Database"))

df_updated.show()
+-----+---------+-----------------+
|db_id|  db_name|          db_type|
+-----+---------+-----------------+
|   12| Teradata|Relation Database|
|   14|Snowflake|Relation Database|
|   15|  Vertica|Relation Database|
|   12| Teradata|Relation Database|
|   22|    Mysql|Relation Database|
+-----+---------+-----------------+

lit() function is used to pass literals i.e. hardcoded values, spark won't take hardcoded values directly.

Example 2: New column db_type_cd is added and updated to literal string "Relation Database" for all the rows.


from pyspark.sql.functions import col,lit
df_updated = df.withColumn("db_type_cd",lit("Relation Database"))

df_updated.show()
+-----+---------+-------+-----------------+
|db_id|  db_name|db_type|       db_type_cd|
+-----+---------+-------+-----------------+
|   12| Teradata|  RDBMS|Relation Database|
|   14|Snowflake|CloudDB|Relation Database|
|   15|  Vertica|  RDBMS|Relation Database|
|   12| Teradata|  RDBMS|Relation Database|
|   22|    Mysql|  RDBMS|Relation Database|
+-----+---------+-------+-----------------+

lit() function is used to pass literals i.e. hardcoded values, spark won't take hardcoded values directly.

Example 3: New column db_name_first_4_char is added and its value is updated using substring function on the db_name column.


from pyspark.sql.functions import col,lit,substring
df_updated = df.withColumn("db_name_first_4_char",substring("db_name",1,4))

df_updated.show()
+-----+---------+-------+--------------------+
|db_id|  db_name|db_type|db_name_first_4_char|
+-----+---------+-------+--------------------+
|   12| Teradata|  RDBMS|                Tera|
|   14|Snowflake|CloudDB|                Snow|
|   15|  Vertica|  RDBMS|                Vert|
|   12| Teradata|  RDBMS|                Tera|
|   22|    Mysql|  RDBMS|                Mysq|
+-----+---------+-------+--------------------+

➠ Add Column using select: select() function can be used with existing columns to add new column to the dataframe. Only downside is that you have to specify all the columns(list can be accessed using df.columns) along with new column.

Example 1: Value of existing column db_type is updated to literal string "Relation Database" for all the rows.


from pyspark.sql.functions import col,lit
df_updated = df.select("db_id", "db_name", lit("Relation Database").alias("db_type"))

df_updated.show()
+-----+---------+-----------------+
|db_id|  db_name|          db_type|
+-----+---------+-----------------+
|   12| Teradata|Relation Database|
|   14|Snowflake|Relation Database|
|   15|  Vertica|Relation Database|
|   12| Teradata|Relation Database|
|   22|    Mysql|Relation Database|
+-----+---------+-----------------+

Example 2: New column db_type_cd is added and updated to literal string "Relation Database" for all the rows.


from pyspark.sql.functions import col,lit
df_updated = df.select("db_id", "db_name", "db_type", lit("Relation Database").alias("db_type_cd"))

df_updated.show()
+-----+---------+-------+-----------------+
|db_id|  db_name|db_type|       db_type_cd|
+-----+---------+-------+-----------------+
|   12| Teradata|  RDBMS|Relation Database|
|   14|Snowflake|CloudDB|Relation Database|
|   15|  Vertica|  RDBMS|Relation Database|
|   12| Teradata|  RDBMS|Relation Database|
|   22|    Mysql|  RDBMS|Relation Database|
+-----+---------+-------+-----------------+



#same example as above but by using list
column_li = df.columns 
column_li.append(lit("Relation Database").alias("db_type_cd"))
df_updated = df.select(column_li)

df_updated.show()
+-----+---------+-------+-----------------+
|db_id|  db_name|db_type|       db_type_cd|
+-----+---------+-------+-----------------+
|   12| Teradata|  RDBMS|Relation Database|
|   14|Snowflake|CloudDB|Relation Database|
|   15|  Vertica|  RDBMS|Relation Database|
|   12| Teradata|  RDBMS|Relation Database|
|   22|    Mysql|  RDBMS|Relation Database|
+-----+---------+-------+-----------------+

lit() function is used to pass literals i.e. hardcoded values, spark won't take hardcoded values directly.

Example 3: New column db_name_first_4_char is added and its value is updated using substring function on the db_name column.


from pyspark.sql.functions import col,lit,substring
df_updated = df.select("db_id", "db_name", "db_type", substring("db_name",1,4).alias("db_name_first_4_char"))

df_updated.show()
+-----+---------+-------+--------------------+
|db_id|  db_name|db_type|db_name_first_4_char|
+-----+---------+-------+--------------------+
|   12| Teradata|  RDBMS|                Tera|
|   14|Snowflake|CloudDB|                Snow|
|   15|  Vertica|  RDBMS|                Vert|
|   12| Teradata|  RDBMS|                Tera|
|   22|    Mysql|  RDBMS|                Mysq|
+-----+---------+-------+--------------------+



#same example as above but by using list
column_li = df.columns
column_li.append(substring("db_name",1,4).alias("db_name_first_4_char"))
df_updated = df.select(column_li)

df_updated.show()
+-----+---------+-------+--------------------+
|db_id|  db_name|db_type|db_name_first_4_char|
+-----+---------+-------+--------------------+
|   12| Teradata|  RDBMS|                Tera|
|   14|Snowflake|CloudDB|                Snow|
|   15|  Vertica|  RDBMS|                Vert|
|   12| Teradata|  RDBMS|                Tera|
|   22|    Mysql|  RDBMS|                Mysq|
+-----+---------+-------+--------------------+

substring() function is used to get a part of string from the db_name column.

Example 4: Adding multiple columns at once using select function.


from pyspark.sql.functions import col,lit,substring
df_updated = df.select("db_id", "db_name", "db_type", substring("db_name",1,4).alias("db_name_first_4_char"), (col("db_id")%10).alias("bucket") )

df_updated.show()
+-----+---------+-------+--------------------+------+
|db_id|  db_name|db_type|db_name_first_4_char|bucket|
+-----+---------+-------+--------------------+------+
|   12| Teradata|  RDBMS|                Tera|   2.0|
|   14|Snowflake|CloudDB|                Snow|   4.0|
|   15|  Vertica|  RDBMS|                Vert|   5.0|
|   12| Teradata|  RDBMS|                Tera|   2.0|
|   22|    Mysql|  RDBMS|                Mysq|   2.0|
+-----+---------+-------+--------------------+------+



#same example as above but by using list
column_li = df.columns 
column_li.append(substring("db_name",1,4).alias("db_name_first_4_char")) #adding 1st Column to list
column_li.append((col("db_id")%10).alias("bucket")) #adding 2nd Column to list
df_updated = df.select(column_li)

df_updated.show()
+-----+---------+-------+--------------------+------+
|db_id|  db_name|db_type|db_name_first_4_char|bucket|
+-----+---------+-------+--------------------+------+
|   12| Teradata|  RDBMS|                Tera|   2.0|
|   14|Snowflake|CloudDB|                Snow|   4.0|
|   15|  Vertica|  RDBMS|                Vert|   5.0|
|   12| Teradata|  RDBMS|                Tera|   2.0|
|   22|    Mysql|  RDBMS|                Mysq|   2.0|
+-----+---------+-------+--------------------+------+

➠ Add Column using other dataframe: Column can be added using other dataframe with the help of outer joins. You can visit dataframe join page to understand more about joins.

Example 1: Columns other_db_name and other_db_type have been added in "df" dataframe using "df_other" dataframe with the help of left outer join.


from pyspark.sql.functions import col,when
df_updated = df.join(df_other, "db_id", "left").select("db_id", df.db_name, df.db_type, df_other.db_name.alias("other_db_name"), df_other.db_type.alias("other_db_type"))

df_updated.show()
+-----+---------+-------+-------------+-------------+
|db_id|  db_name|db_type|other_db_name|other_db_type|
+-----+---------+-------+-------------+-------------+
|   12| Teradata|  RDBMS|         null|         null|
|   14|Snowflake|CloudDB|    Snowflake|        RDBMS|
|   15|  Vertica|  RDBMS|         null|         null|
|   12| Teradata|  RDBMS|         null|         null|
|   22|    Mysql|  RDBMS|        Mysql|        RDBMS|
+-----+---------+-------+-------------+-------------+

This tutorial will explain various approaches with examples on how to add new columns or modify existing columns in a dataframe.

dbmstutorials.com

PySpark: Dataframe Add Columns