Eric Campbell 5 min read Python

PySpark - Update Column Value Conditionally

Conditional column updates are fundamental operations in PySpark, appearing in virtually every data pipeline. Whether you're cleaning messy data, engineering features for machine learning models, or...

Key Insights

  • Use when().otherwise() for most conditional updates—it’s readable, performant, and handles both simple and complex logic through chaining
  • Avoid UDFs for conditional logic unless absolutely necessary; built-in PySpark functions are 5-10x faster due to Catalyst optimizer integration
  • Always use withColumn() to update existing columns rather than creating new DataFrames, as PySpark optimizes these operations into a single pass

Introduction

Conditional column updates are fundamental operations in PySpark, appearing in virtually every data pipeline. Whether you’re cleaning messy data, engineering features for machine learning models, or transforming data during ETL processes, you’ll need to update column values based on specific conditions.

Unlike pandas where you might use boolean indexing or apply(), PySpark requires a different approach due to its distributed nature. The framework provides several methods for conditional updates, each with distinct performance characteristics and use cases.

Let’s start with a sample dataset representing customer transactions:

from pyspark.sql import SparkSession
from pyspark.sql.functions import when, col, expr
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType

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

data = [
    ("C001", 25, "bronze", 150.0),
    ("C002", 45, "silver", 450.0),
    ("C003", 17, "bronze", 80.0),
    ("C004", 52, "gold", 1200.0),
    ("C005", 33, "silver", 320.0),
    ("C006", 19, "bronze", 95.0),
    ("C007", 61, "gold", 2100.0)
]

schema = StructType([
    StructField("customer_id", StringType(), True),
    StructField("age", IntegerType(), True),
    StructField("status", StringType(), True),
    StructField("purchase_amount", DoubleType(), True)
])

df = spark.createDataFrame(data, schema)
df.show()

This creates a DataFrame with customer demographics and transaction data—a typical scenario where conditional updates are needed.

Using when() and otherwise() Functions

The when().otherwise() pattern is your primary tool for conditional updates in PySpark. It mirrors SQL’s CASE WHEN logic but with a more Pythonic syntax.

For a simple condition, let’s update customer status based on age:

from pyspark.sql.functions import when, col

# Update status to 'senior' for customers aged 60+
df_updated = df.withColumn(
    "status",
    when(col("age") >= 60, "senior")
    .otherwise(col("status"))
)

df_updated.show()

The otherwise() clause is crucial—it specifies what happens when the condition isn’t met. Here, we preserve the original status value.

For multiple conditions, chain when() statements:

# Categorize customers into age groups
df_categorized = df.withColumn(
    "age_group",
    when(col("age") < 18, "minor")
    .when((col("age") >= 18) & (col("age") < 30), "young_adult")
    .when((col("age") >= 30) & (col("age") < 50), "adult")
    .when((col("age") >= 50) & (col("age") < 65), "senior_adult")
    .otherwise("senior")
)

df_categorized.select("customer_id", "age", "age_group").show()

Notice the use of & for AND operations and parentheses around each condition—this is mandatory in PySpark. The conditions are evaluated top-to-bottom, and the first matching condition wins.

Here’s a more complex example combining multiple columns:

# Update status based on both age and purchase amount
df_premium = df.withColumn(
    "status",
    when((col("age") >= 50) & (col("purchase_amount") > 1000), "platinum")
    .when((col("purchase_amount") > 500) & (col("purchase_amount") <= 1000), "gold")
    .when(col("purchase_amount") > 200, "silver")
    .otherwise("bronze")
)

df_premium.show()

Using withColumn() for Column Updates

The withColumn() method is how you apply conditional logic to create or update columns. Understanding the distinction between updating existing columns versus creating new ones is important for memory efficiency.

Replacing an existing column (same column name):

# Update purchase_amount with a discount for seniors
df_discounted = df.withColumn(
    "purchase_amount",
    when(col("age") >= 60, col("purchase_amount") * 0.9)
    .otherwise(col("purchase_amount"))
)

df_discounted.show()

Creating a new column based on conditions from multiple columns:

# Create a discount_applied flag
df_with_flag = df.withColumn(
    "discount_applied",
    when((col("age") >= 60) | (col("purchase_amount") > 1000), True)
    .otherwise(False)
)

df_with_flag.show()

Combining mathematical operations with conditional logic:

# Calculate bonus points with conditional multipliers
df_with_points = df.withColumn(
    "bonus_points",
    when(col("status") == "gold", col("purchase_amount") * 0.1)
    .when(col("status") == "silver", col("purchase_amount") * 0.05)
    .otherwise(col("purchase_amount") * 0.02)
)

df_with_points.select("customer_id", "status", "purchase_amount", "bonus_points").show()

You can chain multiple withColumn() calls, but PySpark’s Catalyst optimizer will combine them into a single pass:

df_transformed = df \
    .withColumn("age_group", when(col("age") < 30, "young").otherwise("mature")) \
    .withColumn("high_value", when(col("purchase_amount") > 500, True).otherwise(False)) \
    .withColumn("priority", when((col("age") >= 50) & (col("purchase_amount") > 500), 1).otherwise(0))

df_transformed.show()

Complex Conditional Logic with SQL Expressions

For teams migrating from SQL or dealing with very complex conditions, expr() and selectExpr() allow you to write SQL-style CASE statements directly:

# Using expr() with CASE WHEN
df_case = df.withColumn(
    "customer_tier",
    expr("""
        CASE 
            WHEN age >= 60 AND purchase_amount > 1000 THEN 'VIP Senior'
            WHEN age >= 60 THEN 'Senior'
            WHEN purchase_amount > 1000 THEN 'VIP'
            WHEN purchase_amount > 500 THEN 'Premium'
            ELSE 'Standard'
        END
    """)
)

df_case.select("customer_id", "age", "purchase_amount", "customer_tier").show()

Nested conditions with multiple WHEN clauses:

# Complex business logic using SQL expression
df_complex = df.withColumn(
    "risk_category",
    expr("""
        CASE 
            WHEN age < 18 THEN 'restricted'
            WHEN age >= 18 AND age < 25 AND purchase_amount > 1000 THEN 'high_risk'
            WHEN age >= 25 AND age < 60 AND purchase_amount > 500 THEN 'medium_risk'
            WHEN age >= 60 THEN 'low_risk'
            ELSE 'standard'
        END
    """)
)

df_complex.select("customer_id", "age", "purchase_amount", "risk_category").show()

The equivalent using when().otherwise():

df_equivalent = df.withColumn(
    "risk_category",
    when(col("age") < 18, "restricted")
    .when((col("age") >= 18) & (col("age") < 25) & (col("purchase_amount") > 1000), "high_risk")
    .when((col("age") >= 25) & (col("age") < 60) & (col("purchase_amount") > 500), "medium_risk")
    .when(col("age") >= 60, "low_risk")
    .otherwise("standard")
)

Both approaches produce identical execution plans. Choose based on team familiarity and readability preferences.

Using UDFs for Custom Conditional Logic

User Defined Functions (UDFs) should be your last resort. They bypass Catalyst optimization and serialize/deserialize data between JVM and Python, causing significant performance overhead.

That said, sometimes complex business logic requires UDFs:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# Simple UDF with conditional logic
def categorize_customer(age, amount):
    if age < 18:
        return "restricted"
    elif age >= 60 and amount > 1000:
        return "premium_senior"
    elif amount > 1000:
        return "premium"
    elif amount > 500:
        return "standard_plus"
    else:
        return "standard"

categorize_udf = udf(categorize_customer, StringType())

df_udf = df.withColumn(
    "category",
    categorize_udf(col("age"), col("purchase_amount"))
)

df_udf.show()

UDF with multiple input columns for complex calculations:

def calculate_loyalty_score(age, status, amount):
    base_score = amount * 0.1
    
    if status == "gold":
        multiplier = 1.5
    elif status == "silver":
        multiplier = 1.2
    else:
        multiplier = 1.0
    
    age_bonus = 10 if age >= 60 else 0
    
    return base_score * multiplier + age_bonus

loyalty_udf = udf(calculate_loyalty_score, DoubleType())

df_loyalty = df.withColumn(
    "loyalty_score",
    loyalty_udf(col("age"), col("status"), col("purchase_amount"))
)

df_loyalty.show()

Performance Warning: UDFs are typically 5-10x slower than built-in functions. The above loyalty score calculation should be rewritten using native functions:

# Much faster equivalent without UDF
df_loyalty_optimized = df.withColumn(
    "loyalty_score",
    (col("purchase_amount") * 0.1 * 
     when(col("status") == "gold", 1.5)
     .when(col("status") == "silver", 1.2)
     .otherwise(1.0)) +
    when(col("age") >= 60, 10).otherwise(0)
)

Performance Considerations and Best Practices

Choose built-in functions over UDFs. The Catalyst optimizer can push down predicates, prune columns, and optimize built-in functions. UDFs are black boxes.

Avoid multiple passes over data. Chain withColumn() calls instead of creating intermediate DataFrames:

# Bad - multiple actions
df1 = df.withColumn("col1", when(col("age") > 30, 1).otherwise(0))
df1.show()  # Action 1
df2 = df1.withColumn("col2", when(col("amount") > 500, 1).otherwise(0))
df2.show()  # Action 2

# Good - single pass
df_result = df \
    .withColumn("col1", when(col("age") > 30, 1).otherwise(0)) \
    .withColumn("col2", when(col("amount") > 500, 1).otherwise(0))
df_result.show()  # Single action

Use column pruning. Select only needed columns before complex transformations:

# Process only necessary columns
df_subset = df.select("customer_id", "age", "purchase_amount")
df_processed = df_subset.withColumn(
    "category",
    when(col("age") >= 60, "senior").otherwise("regular")
)

Cache strategically. If you’re applying multiple conditional transformations on the same base DataFrame:

df.cache()
df_transform1 = df.withColumn("col1", when(...).otherwise(...))
df_transform2 = df.withColumn("col2", when(...).otherwise(...))
# Both use cached df

Conclusion

For most conditional updates, stick with when().otherwise()—it’s readable, maintainable, and performs excellently. Use expr() with CASE WHEN only when migrating SQL code or when the SQL syntax is genuinely clearer for complex nested conditions.

Reserve UDFs for truly custom logic that can’t be expressed with built-in functions, and always benchmark against native alternatives. The performance penalty is real and significant.

Here’s your decision matrix:

  • Simple conditions (1-3 branches): when().otherwise()
  • Complex nested logic with SQL background: expr() with CASE WHEN
  • Multiple column transformations: Chained withColumn() calls
  • Irreducible custom business logic: UDFs as last resort

Master these patterns, and you’ll handle 99% of conditional update scenarios efficiently in your PySpark pipelines.

Liked this? There's more.

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