Eric Campbell 5 min read Python

PySpark - Select All Columns Except One

When working with PySpark DataFrames, you'll frequently encounter situations where you need to select all columns except one or a few specific ones. This is a common pattern in data engineering...

Key Insights

  • The drop() method is the simplest and most readable way to exclude columns from a PySpark DataFrame, supporting both single and multiple column removal in one call
  • List comprehension with df.columns provides dynamic column filtering when you need programmatic control over which columns to exclude based on patterns or conditions
  • All column exclusion methods have similar performance characteristics, so choose based on readability and maintainability rather than micro-optimizations

Introduction

When working with PySpark DataFrames, you’ll frequently encounter situations where you need to select all columns except one or a few specific ones. This is a common pattern in data engineering workflows—perhaps you’re dropping a temporary ID column after a join, removing sensitive personally identifiable information (PII) before sharing data with another team, or eliminating redundant columns that bloat your dataset.

Unlike pandas, which offers the drop() method as the primary approach, PySpark provides multiple ways to exclude columns, each with different trade-offs in terms of readability and flexibility. Let’s explore these methods using a practical example.

Here’s a sample DataFrame we’ll use throughout this article:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder.appName("ColumnSelection").getOrCreate()

data = [
    (1, "Alice", "alice@email.com", "123-456-7890", 50000, "2024-01-15"),
    (2, "Bob", "bob@email.com", "234-567-8901", 60000, "2024-01-16"),
    (3, "Carol", "carol@email.com", "345-678-9012", 55000, "2024-01-17"),
]

columns = ["id", "name", "email", "phone", "salary", "created_at"]
df = spark.createDataFrame(data, columns)

df.show()

This gives us a DataFrame with employee information including some sensitive data we might want to exclude in various scenarios.

Using the drop() Method

The drop() method is the most straightforward and Pythonic way to exclude columns from a PySpark DataFrame. It’s explicit, readable, and handles most use cases elegantly.

To drop a single column:

# Remove the phone column
df_no_phone = df.drop("phone")
df_no_phone.show()

# Result: DataFrame with id, name, email, salary, created_at

The real power of drop() becomes apparent when you need to remove multiple columns. You can pass multiple column names as separate arguments:

# Remove multiple sensitive columns
df_sanitized = df.drop("phone", "email", "salary")
df_sanitized.show()

# Result: DataFrame with only id, name, created_at

You can also chain multiple drop() calls, though this is less efficient and less readable:

# Chaining drop calls (not recommended for multiple columns)
df_chained = df.drop("phone").drop("email").drop("salary")

The chained approach creates intermediate DataFrames for each operation, making it less efficient than dropping all columns in a single call. Always prefer passing multiple arguments to a single drop() call.

One important behavior to note: drop() is forgiving. If you attempt to drop a column that doesn’t exist, PySpark won’t throw an error—it simply returns the DataFrame unchanged:

# This won't raise an error even though 'nonexistent' doesn't exist
df_safe = df.drop("nonexistent", "phone")

This can be both a feature and a bug. It’s useful for defensive programming but can mask typos in column names.

Using Column Selection with List Comprehension

When you need more dynamic control over which columns to exclude—perhaps based on a pattern, a configuration file, or runtime conditions—list comprehension with column filtering is your best friend.

The basic pattern involves getting all column names with df.columns, filtering them, and then selecting the filtered list:

# Exclude a single column using list comprehension
columns_to_keep = [c for c in df.columns if c != "phone"]
df_filtered = df.select(columns_to_keep)
df_filtered.show()

This approach shines when you need to exclude multiple columns:

# Exclude multiple columns
exclude_columns = ["phone", "email", "salary"]
columns_to_keep = [c for c in df.columns if c not in exclude_columns]
df_filtered = df.select(columns_to_keep)
df_filtered.show()

You can also use more sophisticated filtering logic. For example, excluding all columns that start with a certain prefix:

# Exclude all columns starting with 'created'
columns_to_keep = [c for c in df.columns if not c.startswith("created")]
df_filtered = df.select(columns_to_keep)

Or exclude columns matching a pattern:

import re

# Exclude columns containing contact information (email or phone)
columns_to_keep = [c for c in df.columns if not re.search(r'email|phone', c)]
df_filtered = df.select(columns_to_keep)

This method is particularly valuable in production pipelines where the columns to exclude might come from a configuration file or database:

# Columns to exclude from configuration
SENSITIVE_COLUMNS = ["email", "phone", "salary"]

def sanitize_dataframe(df, exclude_cols):
    columns_to_keep = [c for c in df.columns if c not in exclude_cols]
    return df.select(columns_to_keep)

df_sanitized = sanitize_dataframe(df, SENSITIVE_COLUMNS)

Using Column Expressions with Unpacking

PySpark’s select() method accepts column expressions, and Python’s unpacking operator (*) lets you pass a list of columns dynamically. This creates another elegant pattern for column exclusion:

# Using unpacking with filtered columns
exclude_cols = ["phone", "email"]
selected_cols = [col(c) for c in df.columns if c not in exclude_cols]
df_selected = df.select(*selected_cols)
df_selected.show()

You can also use string column names directly without wrapping them in col():

# Simpler version without col()
exclude_cols = ["phone", "email"]
selected_cols = [c for c in df.columns if c not in exclude_cols]
df_selected = df.select(*selected_cols)

This approach is functionally equivalent to the list comprehension method shown earlier, but some developers prefer the explicit unpacking syntax for clarity.

When dealing with column name conflicts or special characters, using col() provides additional safety:

# Handling columns with special characters or spaces
from pyspark.sql.functions import col

exclude_cols = ["phone"]
selected_cols = [col(c) for c in df.columns if c not in exclude_cols]
df_selected = df.select(*selected_cols)

Performance Considerations and Best Practices

All the methods discussed above are lazy operations in PySpark—they don’t actually execute until you call an action like show(), collect(), or write(). The performance differences between these approaches are negligible for most use cases.

Let’s verify this by examining the execution plans:

# Compare execution plans
exclude_cols = ["phone", "email"]

# Method 1: drop()
df1 = df.drop("phone", "email")

# Method 2: list comprehension with select()
columns_to_keep = [c for c in df.columns if c not in exclude_cols]
df2 = df.select(columns_to_keep)

# Both produce identical execution plans
print("Method 1 (drop):")
df1.explain()

print("\nMethod 2 (select with filtering):")
df2.explain()

Both approaches produce the same logical plan—they’re just different ways of expressing the same transformation.

Best practices to follow:

  1. Use drop() for static column removal: When you know exactly which columns to exclude at development time, drop() is the clearest choice.

  2. Use list comprehension for dynamic filtering: When column exclusion depends on runtime conditions, patterns, or configuration, list comprehension provides the necessary flexibility.

  3. Avoid chaining multiple drop() calls: Always pass all columns to exclude in a single drop() call for better readability.

  4. Handle missing columns explicitly: If you need strict validation, check column existence before dropping:

def safe_drop(df, columns_to_drop):
    existing_cols = set(df.columns)
    invalid_cols = set(columns_to_drop) - existing_cols
    
    if invalid_cols:
        raise ValueError(f"Columns not found: {invalid_cols}")
    
    return df.drop(*columns_to_drop)

Real-World Use Cases

Let’s explore some practical scenarios where column exclusion is essential.

Removing timestamp columns before joins: When joining DataFrames on business keys, you often need to exclude audit columns that would create duplicates:

# Remove timestamp before joining to avoid column conflicts
df1 = orders_df.drop("created_at", "updated_at")
df2 = customers_df.drop("created_at", "updated_at")

result = df1.join(df2, "customer_id", "inner")

Excluding PII data for analytics: When preparing data for analysts who shouldn’t access sensitive information:

PII_COLUMNS = ["email", "phone", "ssn", "address"]

def prepare_analytics_data(df):
    # Remove PII columns
    safe_columns = [c for c in df.columns if c not in PII_COLUMNS]
    return df.select(safe_columns)

analytics_df = prepare_analytics_data(df)

Dropping intermediate calculation columns: After complex transformations, you may have temporary columns that shouldn’t appear in the final output:

from pyspark.sql.functions import col, when

# Add temporary calculation columns
df_with_temp = df.withColumn("salary_bracket", 
    when(col("salary") < 55000, "low")
    .when(col("salary") < 65000, "medium")
    .otherwise("high")
).withColumn("temp_calculation", col("salary") * 1.1)

# Drop temporary columns before final output
final_df = df_with_temp.drop("temp_calculation")

These patterns form the foundation of clean, maintainable PySpark code. Choose the method that best fits your use case: drop() for simplicity and clarity, list comprehension for dynamic filtering, and always prioritize readability over premature optimization.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.