How to remove DataFrame columns in PySpark Azure Databricks?

Are you looking to find how to remove the column of PySpark DataFrame into Azure Databricks cloud or maybe you are looking for a solution, to drop the DataFrame column in PySpark Databricks using the drop method? If you are looking for any of these problem solutions, then you have landed on the correct page. I will also show you what to do and how to use PySpark to remove the column name of DataFrame in Azure Databricks. I will explain it with a practical example. So don’t waste time let’s start with a step-by-step guide to understanding how to drop columns in a PySpark DataFrame.

In this blog, I will teach you the following with practical examples:

  • Syntax of drop()
  • Removing single column
  • Removing multiple columns

drop() method used to remove column of the DataFrame in PySpark Azure Databricks.

Syntax: dataframe_name.drop(column_names)

What is the syntax of the drop() function in PySpark Azure Databricks?

The syntax is as follows:

dataframe_name.drop(*cols)
Parameter NameRequiredDescription
cols (str)YesName of the column, or the Column to drop.
Table 1: drop() Method in PySpark Databricks Parameter list with Details

Apache Spark Official documentation link: drop()

Create a simple DataFrame

Gentle reminder:

In Databricks,

  • sparkSession made available as spark
  • sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

data = [
    (1,"Halley","Sales","https://diigo.com"),
    (2,"Leoine","HR","https://nbcnews.com"),
    (3,"Hyatt","HR","http://youtu.be"),
    (4,"Cathie","Sales","http://tuttocitta.it"),
    (5,"Barnebas","HR","https://e-recht24.de")
]

df = spark.createDataFrame(data, schema = ["id","name","department","url"])
df.printSchema()
df.show(truncate=False)

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- url: string (nullable = true)

+---+--------+----------+--------------------+
|id |name    |department|url                 |
+---+--------+----------+--------------------+
|1  |Halley  |Sales     |https://diigo.com   |
|2  |Leoine  |HR        |https://nbcnews.com |
|3  |Hyatt   |HR        |http://youtu.be     |
|4  |Cathie  |Sales     |http://tuttocitta.it|
|5  |Barnebas|HR        |https://e-recht24.de|
+---+--------+----------+--------------------+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

# replace the file_paths with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- url: string (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

How to remove a single column in PySpark Azure Databricks using the drop() function?

By providing the column name to the drop() function, you can remove or drop the column name from the DataFrame. This generates a new DataFrame with the chosen columns because DataFrame is immutable. To remove the column of a Dataframe, use the drop() function.

Examples:

df.drop("url"').show()

"""
Output:

+---+--------+----------+
| id|    name|department|
+---+--------+----------+
|  1|  Halley|     Sales|
|  2|  Leoine|        HR|
|  3|   Hyatt|        HR|
|  4|  Cathie|     Sales|
|  5|Barnebas|        HR|
+---+--------+----------+

"""

How to remove multiple columns in PySpark Azure Databricks using the drop() function?

To remove multiple columns from a Dataframe, pass the column name as a string to the drop() function.

Example:

df.drop("department", "url").show()

"""
Output:

+---+--------+
| id|    name|
+---+--------+
|  1|  Halley|
|  2|  Leoine|
|  3|   Hyatt|
|  4|  Cathie|
|  5|Barnebas|
+---+--------+

"""

I have attached the complete code used in this blog in notebook format to this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

When should you use the PySpark drop() function in Azure Databricks?

You can use the drop() function to remove one or more columns from a DataFrame. The drop() function creates a new dataframe from the existing dataframe by removing the specified columns.

Real World Use Case Scenarios for PySpark DataFrame drop() function in Azure Databricks?

  • Out of 100 columns, I want to keep only relevant columns in the DataFrame.
  • Want to remove an extra column from DataFrame to reduce the total size of the data to be processed.

What are the alternatives of the drop() function in PySpark Azure Databricks?

Use the select() function to get specific columns instead of dropping them.

Final Thoughts

In this article, we have learned about the PySpark drop() method to remove the columns of DataFrame in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

Arud Seka Berne S

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.