PySpark Drop Columns – Eliminate Unwanted Columns in PySpark DataFrame with Ease
Detailed blog post on using PySpark's Drop() function to remove columns from a DataFrame, explore various use cases to understand its versatility and importance in data manipulation
Welcome to this detailed blog post on using PySpark’s Drop() function to remove columns from a DataFrame. Lets delve into the mechanics of the Drop() function and explore various use cases to understand its versatility and importance in data manipulation.
This post is a perfect starting point for those looking to expand their understanding of PySpark and improve their data wrangling skills.
Creating a DataFrame
Before we dive into the Drop() function, let’s create a DataFrame to work with. In this example, we will create a simple DataFrame with four columns: “name”, “age”, “city”, and “gender.”
import findspark
findspark.init()
from pyspark.sql import SparkSession
spark = SparkSession.builder \
.appName("PySpark Drop Column Example") \
.getOrCreate()
data = [("Alice", 30, "New York", "F"),
("Bob", 28, "San Francisco", "M"),
("Cathy", 29, "Los Angeles", "F"),
("David", 32, "Chicago", "M")]
columns = ["name", "age", "city", "gender"]
df = spark.createDataFrame(data, columns)
df.show()
+-----+---+-------------+------+
| name|age| city|gender|
+-----+---+-------------+------+
|Alice| 30| New York| F|
| Bob| 28|San Francisco| M|
|Cathy| 29| Los Angeles| F|
|David| 32| Chicago| M|
+-----+---+-------------+------+
Different ways to drop columns in PySpark DataFrame
- Dropping a Single Column
- Dropping Multiple Columns
- Dropping Columns Conditionally
- Dropping Columns Using Regex Pattern
1. Dropping a Single Column
The Drop() function can be used to remove a single column from a DataFrame. The syntax is as follows
df = df.drop("gender")
df.show()
+-----+-------------+
| name| city|
+-----+-------------+
|Alice| New York|
| Bob|San Francisco|
|Cathy| Los Angeles|
|David| Chicago|
+-----+-------------+
2. Dropping Multiple Columns:
You can also use the Drop() function to remove multiple columns from a DataFrame. Simply pass a list of column names to the function
For example, let’s remove both “age” and “gender” columns
df = df.drop("age", "gender")
df.show()
+-----+-------------+
| name| city|
+-----+-------------+
|Alice| New York|
| Bob|San Francisco|
|Cathy| Los Angeles|
|David| Chicago|
+-----+-------------+
Alternatively, you can use a list of column names
columns_to_drop = ["age", "gender"]
df = df.drop(*columns_to_drop)
df.show()
+-----+-------------+
| name| city|
+-----+-------------+
|Alice| New York|
| Bob|San Francisco|
|Cathy| Los Angeles|
|David| Chicago|
+-----+-------------+
3. Dropping Columns Conditionally
You might want to drop columns based on a specific condition. You can use the Drop() function in combination with the “if” statement to achieve this
if "gender" in df.columns:
df = df.drop("gender")
df.show()
+-----+---+-------------+
| name|age| city|
+-----+---+-------------+
|Alice| 30| New York|
| Bob| 28|San Francisco|
|Cathy| 29| Los Angeles|
|David| 32| Chicago|
+-----+---+-------------+
4. Dropping Columns Using Regex Pattern
You can use the “drop()” function in combination with a regular expression (regex) pattern to drop multiple columns matching the pattern.
from pyspark.sql.functions import col
import re
regex_pattern = "gender|age"
df = df.select([col(c) for c in df.columns if not re.match(regex_pattern, c)])
df.show()
+-----+-------------+
| name| city|
+-----+-------------+
|Alice| New York|
| Bob|San Francisco|
|Cathy| Los Angeles|
|David| Chicago|
+-----+-------------+
Conclusion
In this blog post, we learned about the PySpark Drop() function and its various use cases. We explored how to remove single and multiple columns, drop columns conditionally, and remove columns using a regex pattern.
With a solid understanding of the PySpark Drop() function, you can now effectively manipulate your data to suit your needs.
Build a strong Python foundation with hands-on exercises designed for aspiring Data Scientists and AI/ML Engineers.
Start Free Course →
Related Articles
