Databricks Notebook Optimization for Large-Scale ETL

Databricks Notebook Optimization for Large-Scale ETL

By a Principal Data Engineer & Databricks Specialist

1. Introduction

I’ve spent the better part of two decades in data engineering, with a significant portion of my early career mastering DataStage to build enterprise-scale ETL. The move to Databricks over the last 10 years has been transformative, but it hasn't been without its challenges. Migrating large, complex ETL workloads from a visual, stage-based tool like DataStage to code-first Databricks notebooks is more than a simple lift-and-shift; it's a paradigm shift.

The primary difference lies in control and responsibility. DataStage abstracts away much of the underlying parallelism and resource management. In Databricks, the power of Spark gives you near-infinite flexibility, but it also hands you the responsibility for managing that power. An unoptimized notebook in a development environment might run fine on a small dataset, but in production, with terabytes of data, that same code can bring a cluster to its knees, blow through budgets, and miss critical SLAs.

This guide is for my fellow senior engineers, architects, and technical leads. It’s a collection of hard-won lessons from a decade of migrating and optimizing large-scale ETL on Databricks. We will skip the marketing and dive straight into the technical details of building high-performance, cost-efficient, and reliable notebook-based ETL.

2. Common Challenges in Large-Scale Notebooks

When I'm brought in to review a struggling Databricks project, the symptoms are almost always the same. Understanding these is the first step to diagnosis.

• Long Execution Times: Jobs that should take minutes take hours. This is the most obvious and painful symptom, directly impacting data freshness and SLAs.
• Skewed Data & Inefficient Joins: One or two tasks in a Spark stage run for hours while the other 198 finish in seconds. This is a classic sign of data skew, where the partitioning key is unevenly distributed, overwhelming specific executors.
• Unnecessary Shuffles: The Spark UI's DAG visualization looks like a tangled web. A "shuffle" is Spark moving data between executors across the network. It's the most expensive operation in distributed computing. Excessive or repeated shuffles are a primary cause of poor performance.
• Resource Overuse & High Cost: Clusters are provisioned with hundreds of cores, running for hours, yet CPU utilization is low. This points to I/O-bound operations or inefficient code that isn't leveraging the available compute, leading to massive and unnecessary cloud bills.

3. Code Optimization Techniques

Your code is the foundation. A poorly configured cluster cannot fix poorly written Spark code.

Refactoring DataStage Logic into Efficient PySpark/DataFrame Code

The mindset must shift from visual stages to DataFrame transformations.

• A DataStage Transformer stage with 50 derivations becomes a series of .withColumn() calls.
• A Filter stage is a .filter() or .where() clause.
• A Join stage is a .join() operation.

The key is to chain these transformations together. Spark's Catalyst Optimizer will analyze the entire chain and create an optimized physical execution plan.

Anti-Pattern: Executing an action (like .show() or .count()) after every minor transformation. This forces Spark to materialize the results at each step, breaking the optimization chain.

    # Anti-Pattern: Breaking the optimization chain
    df_source = spark.read.format("delta").load("/path/to/source")
    df_filtered = df_source.filter("status = 'active'")
    print("Filtered count:", df_filtered.count()) # Action #1 - Triggers a job

    df_transformed = df_filtered.withColumn("new_col", F.upper("name"))
    print("Transformed count:", df_transformed.count()) # Action #2 - Triggers another job

    # Better Pattern: Chain transformations, one action at the end
    df_final = (
        spark.read.format("delta").load("/path/to/source")
        .filter("status = 'active'")
        .withColumn("new_col", F.upper("name"))
    )
    df_final.write.format("delta").save("/path/to/target") # One action, one optimized job

Using Spark SQL Where Appropriate

For engineers coming from a heavy SQL background, Spark SQL can be more intuitive and, in many cases, just as performant as the DataFrame API. The Catalyst Optimizer processes both identically. I often find that complex business logic involving multiple joins and aggregations is cleaner and easier to maintain in SQL.

    -- This can be much more readable than a series of .join() and .groupBy() calls
    CREATE OR REPLACE TEMP VIEW final_agg AS
    SELECT
      t.customer_id,
      d.region,
      COUNT(DISTINCT t.order_id) as order_count,
      SUM(t.order_value) as total_spend
    FROM raw_transactions t
    JOIN customer_demographics d ON t.customer_id = d.customer_id
    WHERE t.order_date >= '2023-01-01'
    GROUP BY 1, 2

Avoiding Wide Transformations and Unnecessary Caching

• Narrow Transformations (select, filter): Operate on a single partition; no data movement is required. They are cheap.
• Wide Transformations (groupBy, join, repartition): Require data to be shuffled across the network. They are expensive.

Your goal is to minimize the number and scope of wide transformations. For example, filter your data before you join it. This reduces the amount of data that needs to be shuffled.

Optimizing UDFs and Leveraging Vectorized Operations

Standard Python User-Defined Functions (UDFs) are a performance killer. For each row, Spark has to:
1. Serialize the data from the JVM.
2. Send it to a Python process.
3. Execute the Python code.
4. Serialize the result back.
5. Deserialize it in the JVM.

This overhead is immense on large datasets.

The Hierarchy of Function Performance:
1. Use built-in Spark functions first. They run entirely within the optimized JVM environment.
2. If you must use custom logic, use Pandas UDFs (Vectorized UDFs). They operate on Pandas Series/DataFrames in batches, dramatically reducing the serialization overhead.
3. Avoid row-at-a-time Python UDFs for any performance-critical path.

    from pyspark.sql.functions import pandas_udf
    import pandas as pd

    # Slow: Row-at-a-time Python UDF
    def upper_case_udf(s: str) -> str:
        return s.upper()
    my_udf = F.udf(upper_case_udf, StringType())
    df_slow = df.withColumn("name_upper", my_udf(F.col("name")))


    # Fast: Vectorized Pandas UDF
    @pandas_udf("string")
    def pandas_upper_udf(s: pd.Series) -> pd.Series:
        return s.str.upper()
    df_fast = df.withColumn("name_upper", pandas_upper_udf(F.col("name")))

4. Data Partitioning & Layout

How your data is physically stored on disk is as important as your code.

Repartitioning and Coalescing Strategies

• df.repartition(N): Triggers a full shuffle to create exactly N partitions of roughly equal size. Use this before a write to control output file size or before a join to pre-partition data on the join key.
• df.coalesce(M): Reduces the number of partitions to M by combining existing partitions. It avoids a full shuffle but can lead to unevenly sized partitions. Use this only to reduce a large number of partitions to a smaller one (e.g., from 2000 to 200) right before writing.

Rule of Thumb: Aim for partition sizes between 128MB and 1GB. Too many small files create metadata overhead for the driver; too few large files reduce parallelism.

Handling Skewed Datasets

Data skew is the silent killer of Spark jobs. If you join transactions and users on user_id, and one user_id (e.g., a "guest" or test user) has billions of records, one executor gets all that data and chokes.

Solution: Salting. Add a random component to the join key to distribute the skewed data across multiple executors.

    from pyspark.sql.functions import concat, lit, col, floor, rand

    # Assume 'df_large' is skewed on 'join_key' for a specific value 'SKEWED_VALUE'
    SALT_FACTOR = 10

    # Salt the large DataFrame
    df_large_salted = df_large.withColumn(
        "salted_join_key",
        F.when(
            col("join_key") == "SKEWED_VALUE",
            concat(col("join_key"), lit("_"), floor(rand() * SALT_FACTOR))
        ).otherwise(col("join_key"))
    )

    # Explode the small DataFrame to match all possible salts
    df_small_salted = df_small.crossJoin(
        spark.range(SALT_FACTOR).withColumnRenamed("id", "salt")
    ).withColumn(
        "salted_join_key",
        concat(col("join_key"), lit("_"), col("salt"))
    )

    # Now join on the salted key
    result = df_large_salted.join(df_small_salted, "salted_join_key")

Optimizing File Formats (Parquet, Delta)

• Parquet: The default choice for columnar storage. It enables predicate pushdown (skipping read for files not matching a WHERE clause) and column pruning (only reading the columns you select).
• Delta Lake: The gold standard for ETL on Databricks. It's Parquet plus a transaction log. This gives you ACID transactions, time travel, and, critically, performance optimizations.
  • OPTIMIZE: Compacts small files into larger, optimally sized ones. Run this regularly on tables that receive many small writes.
  • Z-ORDER BY (col1, col2): Colocates related data in the same set of files. If you frequently filter by col1 and col2, Z-Ordering can dramatically speed up reads by reducing the number of files scanned.

5. Cluster & Resource Optimization

No amount of code optimization can save you if your cluster is misconfigured.

• Right-Sizing Clusters: Don't guess. Run your workload on a baseline cluster, then use the Spark UI and cluster metrics to analyze resource usage.
  • Is CPU maxed out? You're compute-bound. Consider compute-optimized instances or more cores.
  • Is the cluster sitting idle waiting on I/O? You're I/O bound. Check for small file problems or inefficient reads. Maybe you need storage-optimized instances.
  • Are you getting OOM (Out of Memory) errors? You're memory-bound. Increase executor memory or use memory-optimized instances.
• Auto-Scaling: Always enable auto-scaling for ETL workloads. This allows the cluster to acquire workers when shuffle-heavy stages begin and release them during less intensive stages. It's a primary lever for cost control.
• Instance Selection: Use the right tool for the job.
  • Memory-Optimized (r-series on AWS, E-series on Azure): Best for workloads with large shuffles, caches, and aggregations. My default starting point for unknown, heavy ETL.
  • Compute-Optimized (c-series on AWS, F-series on Azure): Good for CPU-intensive tasks like complex string manipulations, ML inference, or when code is highly optimized.
• Memory and CPU Tuning: The two most important settings are spark.executor.memory and spark.executor.cores. A common starting point is to assign 4-5 cores per executor. This balances parallelism with I/O throughput. Don't over-allocate memory; leave enough for OS overhead (approx. 10%).

6. Workflow & Notebook Management

Monolithic, 1000-line notebooks are unmaintainable and impossible to debug.

• Modular Notebook Design: Break down your ETL process into logical, single-responsibility notebooks.
  • 01_ingest_data
  • 02_clean_and_transform_customers
  • 03_aggregate_transactions
  • 04_join_and_load_final_table
• Using Databricks Workflows: Use the Workflows UI to orchestrate these notebooks. This allows you to define dependencies and run tasks in parallel. For instance, 02_..._customers and 03_..._transactions can run simultaneously, significantly reducing the total wall-clock time. Pass parameters between notebooks using dbutils.notebook.exit() and task values.
• Dependency Management: Define clear dependencies in your workflow. This creates a Directed Acyclic Graph (DAG) that is easy to visualize, debug, and maintain.

7. Caching & Persistence Strategies

Caching is a powerful tool but is often misused.

• When to Cache (.cache() or .persist()): Cache a DataFrame only when you will access it multiple times in the same job. A classic example is a dimension table that is joined to multiple fact tables within the same run.
• The Pitfall: Caching is lazy. The first action on a cached DataFrame will be slow as it computes and stores the data in memory/disk. Subsequent actions are fast. If you only use a DataFrame once, caching it actually slows down your job. It also consumes memory, which can lead to OOM errors or spilling to disk, negating the benefit.
• Using Delta Tables as an Alternative: A more robust pattern, especially for very large intermediate DataFrames, is to write them to a temporary Delta table. This is called "strategic checkpointing." It's slower than in-memory caching but is fault-tolerant (the data survives if an executor dies) and frees up memory for downstream computations.

8. Monitoring & Observability

You can't optimize what you can't see.

• Spark UI: Your Best Friend. This is non-negotiable.
  • Jobs/Stages Tab: Look for long-running stages. Click into a stage to see its DAG. A wide, complex stage DAG often indicates a performance bottleneck (e.g., a massive shuffle). Identify which tasks are slow (stragglers) to diagnose skew.
  • SQL/DataFrame Tab: Select your query and examine the Physical Plan. This shows you exactly how Spark will execute your code. Look for operations like Exchange (a shuffle) and BroadcastHashJoin vs. SortMergeJoin.
  • Executors Tab: Look at the Shuffle Read/Write and Input/Output columns. Wildly different values across executors confirm data skew.
• Databricks Performance Dashboards: Use the built-in cluster metrics to get a high-level view of CPU, Memory, and Network I/O over the life of the job. A pattern of spiky CPU followed by long flatlines often indicates an I/O bottleneck.
• Metrics-Driven Iterative Tuning: Optimization is a cycle.
  1. Establish a baseline runtime.
  2. Use the Spark UI to form a hypothesis (e.g., "This join is skewed").
  3. Apply one specific optimization (e.g., implement salting).
  4. Re-run the job and measure the new runtime and resource usage.
  5. Repeat.

9. Cost Optimization

Performance and cost are two sides of the same coin. A faster job is almost always a cheaper job.

• Reducing Cluster Idle Time: Use auto-terminating clusters aggressively. A job cluster should terminate within minutes of completion. For ad-hoc queries, set a low inactivity period (e.g., 30-60 minutes).
• Optimizing Execution Schedules: Challenge business requirements. Does this hourly job really need to run every hour, or can it run every 4 hours? Aligning the job schedule with the actual data refresh cadence and business need is a huge cost lever.
• Spot Instances: For non-critical, fault-tolerant workloads, using spot instances for worker nodes can reduce compute costs by up to 90%. Use a hybrid strategy: a reliable on-demand instance for the driver node and spot instances for the auto-scaling workers. Databricks can automatically try to re-acquire spot instances if they are reclaimed.
• Cluster Pools: For jobs that run frequently and need to start quickly, pools of pre-warmed, idle instances can eliminate cluster start-up time. This is a trade-off: you pay a small amount to keep the instances in the pool, but your jobs start instantly. It's great for latency-sensitive interactive workloads or frequent, short jobs.

10. Real-World Examples & Lessons Learned

Lesson 1: The Skewed Sales Aggregation

• Problem: A daily job to aggregate sales by store_id was taking 4+ hours. The Spark UI showed one task in the final aggregation stage running for 3.9 hours, while all others finished in 2 minutes.
• Diagnosis: Investigation revealed a store_id = -1 used for "online sales" which accounted for 80% of the total transaction volume. The groupBy("store_id") was sending all this data to a single executor.
• Optimization: We refactored the code. We first filtered for store_id = -1, aggregated it separately (no groupBy needed), and then processed the remaining well-distributed store data in parallel. Finally, we unioned the two results.
• Result: The job runtime dropped from 4+ hours to 15 minutes.

Lesson 2: The "Death by a Thousand Files" Ingest

• Problem: A downstream dashboard that read from a Delta table was timing out. The table was produced by an hourly streaming notebook that used .writeStream.trigger(availableNow=True).
• Diagnosis: Looking at the Delta table's storage, we found it contained hundreds of thousands of tiny Kilobyte-sized Parquet files. Each query against the table had to list all these files, creating massive metadata overhead on the driver before any data was even read.
• Optimization: We added a final step to the workflow: a separate notebook that ran OPTIMIZE my_delta_table every 6 hours. We also added Z-ORDER BY (event_date) as it was the most common filter predicate.
• Result: The downstream dashboard query time improved from 5+ minutes (timeout) to 8 seconds. The OPTIMIZE job itself took about 10 minutes but saved hours of pain for downstream consumers.

11. Executive Summary / CXO Takeaways

For leadership, the conversation about notebook optimization translates directly to business value.

• Key Risks of Unoptimized ETL:

  • Uncontrolled Cloud Spend: Inefficient jobs are the primary driver of budget overruns. A single bad notebook can cost thousands of dollars per day.
  • Missed SLAs & Stale Data: Slow jobs mean decision-makers are working with outdated information, eroding trust in the data platform.
  • Brittle, Unreliable Pipelines: Unoptimized jobs are often fragile. A small increase in data volume can cause them to fail entirely, requiring constant, reactive firefighting from the engineering team.

• Key Benefits of a Proactive Optimization Strategy:

  • Lower Total Cost of Ownership (TCO): Optimized jobs use fewer resources for less time, directly reducing the monthly cloud bill.
  • Predictable Performance & Reliability: Well-tuned pipelines perform consistently, meeting SLAs and building business confidence.
  • Increased Engineering Productivity: Engineers can focus on building new value instead of debugging performance issues in production.

• Recommendations for Governance & Management:

  1. Establish a Center of Excellence (CoE): Create a small, dedicated team of experts to set best practices, conduct code reviews for critical pipelines, and train other engineers.
  2. Implement Automated Guardrails: Use cluster policies to enforce tagging (for cost attribution), limit max cluster sizes, and set mandatory auto-termination policies.
  3. Mandate Performance Reviews: For any ETL job deemed "critical," a performance and cost review using the Spark UI should be a mandatory part of the pre-deployment checklist.
  4. Invest in Observability: Make sure your teams know how to use the Spark UI and Databricks monitoring tools. You cannot manage costs or performance that you cannot see.

By embracing these principles, we transform Databricks notebooks from a potential liability into a powerful, efficient, and cost-effective engine for large-scale ETL. The platform provides the tools; it is our engineering discipline that ensures they are used wisely.