Drop One or More Columns From Pyspark DataFrame

In this tutorial, we will look at how to drop one or more columns from a Pyspark dataframe with the help of some examples.

How to drop Pyspark dataframe columns?

You can use the Pyspark drop() function to drop one or more columns from a Pyspark dataframe. Pass the column (or columns) you want to drop as arguments to the function. The following is the syntax –

# drop column from dataframe
df.drop("column1", "column2", ...)

It returns a Pyspark dataframe resulting from removing the passed column(s).

Examples

Let’s look at some examples of removing columns from a Pypsark dataframe. First, we will create a Pyspark dataframe that we’ll be using throughout this tutorial.

# import the pyspark module
import pyspark
  
# import the  sparksession class  from pyspark.sql
from pyspark.sql import SparkSession
  
# create an app from SparkSession class
spark = SparkSession.builder.appName('datascience_parichay').getOrCreate()

# data of competition participants
data = [["Tim", 23, "Germany", "A"],
      ["Viraj", 22, "India", "A"],
      ["Emma", 24, "USA", "B"],
      ["Jack", 28, "USA", "B"],
      ["Max", 25, "Germany", "A"]]

# create a Pyspark dataframe using the above data
df = spark.createDataFrame(data, ["Name", "Age", "Country", "Team"])

# display 
df.show()

Output:

+-----+---+-------+----+
| Name|Age|Country|Team|
+-----+---+-------+----+
|  Tim| 23|Germany|   A|
|Viraj| 22|  India|   A|
| Emma| 24|    USA|   B|
| Jack| 28|    USA|   B|
|  Max| 25|Germany|   A|
+-----+---+-------+----+

We now have a dataframe containing the name, age, country, and team information of some students participating in a case-study competition.

Drop a column from Pyspark dataframe

Let’s drop the “Age” column from the above dataframe. For this, apply the drop() function to the dataframe and pass “Age” as the argument.

# drop "Age" column
df.drop("Age").show()

Output:

+-----+-------+----+
| Name|Country|Team|
+-----+-------+----+
|  Tim|Germany|   A|
|Viraj|  India|   A|
| Emma|    USA|   B|
| Jack|    USA|   B|
|  Max|Germany|   A|
+-----+-------+----+

The resulting dataframe has the “Age” column removed. Note that the drop() function doesn’t modify the original dataframe in place. To modify the original dataframe, assign the resulting dataframe from the drop() function to the original dataframe variable.

# drop "Age" column
df = df.drop("Age")
# display the dataframe
df.show()

Output:

+-----+-------+----+
| Name|Country|Team|
+-----+-------+----+
|  Tim|Germany|   A|
|Viraj|  India|   A|
| Emma|    USA|   B|
| Jack|    USA|   B|
|  Max|Germany|   A|
+-----+-------+----+

You can see that the dataframe now doesn’t have the “Age” column.

Drop multiple columns from Pyspark dataframe

You can also use the drop() function to remove more than one column from a Pyspark dataframe. Pass the columns you want to drop as arguments to the drop() function.

For example, let’s drop the columns “Country” and “Team” from the above dataframe.

# drop "Country" and "Team" columns
df.drop("Country", "Team").show()

Output:

+-----+
| Name|
+-----+
|  Tim|
|Viraj|
| Emma|
| Jack|
|  Max|
+-----+

You can see that the resulting dataframe has both columns removed.

You might also be interested in –

• Order PySpark DataFrame using orderBy()
• Filter PySpark DataFrame with where()
• Aggregate Functions in PySpark