Optimizing Databricks Spark jobs using dynamic partition pruning and AQE

Learn how to supercharge your Databricks Spark jobs using Dynamic Partition Pruning (DPP) and Adaptive Query Execution (AQE). This comprehensive guide walks through practical implementations, real-world scenarios, and best practices for optimizing large-scale data processing. Discover how to significantly reduce query execution time and resource usage through intelligent partition handling and runtime optimizations. Perfect for data engineers and architects looking to enhance their Spark job performance in Databricks environments.

GraphQL has a role beyond API Query Language- being the backbone of application Integration
background Coditation

Optimizing Databricks Spark jobs using dynamic partition pruning and AQE

Optimizing Databricks Spark jobs using Dynamic Partition Pruning and AQE

Spark jobs processing large datasets often face performance challenges, particularly when dealing with join operations and partitioned tables. Two powerful features in Databricks - Dynamic Partition Pruning (DPP) and Adaptive Query Execution (AQE) - can significantly improve query performance and reduce resource utilization.

Let's explore these optimization techniques through practical examples and real-world scenarios.

Understanding Dynamic Partition Pruning

Dynamic Partition Pruning enables Spark to skip reading unnecessary partitions during join operations. When a large partitioned fact table joins with a smaller dimension table, DPP determines which partitions are relevant before reading the data.

Consider a retail analytics scenario with these tables: - sales_transactions (1TB, partitioned by date) - store_details (10MB)

Without DPP, querying sales data for specific stores would scan the entire sales_transactions table. With DPP, Spark reads only the required partitions.

Here's a practical example:

-- Create partitioned sales table
CREATE TABLE sales_transactions (
    transaction_id BIGINT,
    store_id INT,
    sale_date DATE,
    amount DECIMAL(10,2)
)
USING DELTA
PARTITIONED BY (sale_date);

-- Create store details table
CREATE TABLE store_details (
    store_id INT,
    region STRING,
    store_type STRING
);

-- Query that benefits from DPP
SELECT 
    s.store_id,
    d.region,
    SUM(s.amount) as total_sales
FROM sales_transactions s
JOIN store_details d ON s.store_id = d.store_id
WHERE d.region = 'WEST'
GROUP BY s.store_id, d.region;

To verify DPP in action:

spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")

# Check execution plan
spark.sql("""
    EXPLAIN EXTENDED
    SELECT s.store_id, d.region, SUM(s.amount) as total_sales
    FROM sales_transactions s
    JOIN store_details d ON s.store_id = d.store_id
    WHERE d.region = 'WEST'
    GROUP BY s.store_id, d.region
""").show(truncate=False)

Implementing Adaptive Query Execution

AQE optimizes query plans during runtime based on accurate statistics. It provides three main optimizations:

  1. Dynamic coalesce shuffle partitions
  2. Dynamic switch join strategies
  3. Dynamic optimize skew joins

Enable AQE in Databricks:

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

Let's implement a scenario handling skewed data:

from pyspark.sql import functions as F

# Create sample skewed dataset
def generate_skewed_data(spark, num_records):
    return (spark.range(num_records)
           .withColumn("key", F.when(F.col("id") % 100 == 0, 1)
           .otherwise(F.col("id")))
           .withColumn("value", F.rand()))

# Generate datasets
large_df = generate_skewed_data(spark, 10000000)
small_df = spark.range(1000).withColumn("value", F.rand())

# Cache dataframes
large_df.cache()
small_df.cache()

# Perform join with AQE
result = large_df.join(small_df, large_df.key == small_df.id)

Monitoring and Tuning

To monitor DPP and AQE effectiveness:

def analyze_query_metrics(query):
    # Enable detailed metrics
    spark.conf.set("spark.sql.execution.metrics.enabled", "true")

    # Execute query and capture metrics
    start_time = time.time()
    result = spark.sql(query)
    result.collect()
    execution_time = time.time() - start_time

    # Get metrics from last query execution
    last_query_metrics = spark.sql("SELECT * FROM sys.last_query_metrics").collect()[0]

    return {
        "execution_time": execution_time,
        "shuffle_read_bytes": last_query_metrics.shuffle_read_bytes,
        "shuffle_write_bytes": last_query_metrics.shuffle_write_bytes,
        "partition_count": last_query_metrics.partition_count
    }

Performance Impact Analysis

Here's a real-world performance comparison from a production environment:

# Test query
test_query = """
SELECT 
    year(s.sale_date) as year,
    d.region,
    SUM(s.amount) as total_sales
FROM sales_transactions s
JOIN store_details d ON s.store_id = d.store_id
WHERE d.region IN ('WEST', 'EAST')
GROUP BY year(s.sale_date), d.region
"""

# Test with different configurations
configurations = [
    {"dpp": False, "aqe": False},
    {"dpp": True, "aqe": False},
    {"dpp": True, "aqe": True}
]

results = {}
for config in configurations:
    spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", 
                  str(config["dpp"]).lower())
    spark.conf.set("spark.sql.adaptive.enabled", 
                  str(config["aqe"]).lower())

    results[f"dpp_{config['dpp']}_aqe_{config['aqe']}"] = analyze_query_metrics(test_query)

Fine-tuning AQE Parameters

Optimize AQE behavior with these configurations:

# Adjust partition size for coalescing
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "64m")
spark.conf.set("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "200")

# Configure skew join handling
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256mb")

# Set advisory partition size
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")

Handling Complex Scenarios

For complex queries involving multiple joins and aggregations:

# Example of a complex analytical query
complex_query = """
WITH daily_sales AS (
    SELECT 
        sale_date,
        store_id,
        SUM(amount) as daily_total
    FROM sales_transactions
    GROUP BY sale_date, store_id
),
store_rankings AS (
    SELECT 
        d.region,
        s.store_id,
        RANK() OVER (PARTITION BY d.region ORDER BY SUM(ds.daily_total) DESC) as rank
    FROM daily_sales ds
    JOIN store_details d ON ds.store_id = d.store_id
    GROUP BY d.region, s.store_id
)
SELECT 
    region,
    store_id,
    rank
FROM store_rankings
WHERE rank <= 10
"""

# Enable all optimizations
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.localShuffleReader.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# Execute and analyze
complex_query_metrics = analyze_query_metrics(complex_query)

Best Practices and Common Pitfalls

  1. Partition Column Selection: 

    # Good: Date-based partitioning for time-series data
spark.sql("""
    CREATE TABLE good_partitioning (
        id BIGINT,
        timestamp TIMESTAMP,
        value DOUBLE
    )
    USING DELTA
    PARTITIONED BY (date_trunc('day', timestamp))
""")

# Bad: Too fine-grained partitioning
spark.sql("""
    CREATE TABLE bad_partitioning (
        id BIGINT,
        timestamp TIMESTAMP,
        value DOUBLE
    )
    USING DELTA
    PARTITIONED BY (timestamp)
""")

  2. Handling Data Skew: 

    # Add salting for skewed columns
def salt_skewed_column(df, column, num_salts=10):
    return df.withColumn(
        "salted_key",
        F.concat(
            F.col(column),
            F.lit("_"),
            (F.rand() * num_salts).cast("int").cast("string")
        )
    )

# Apply to skewed dataset
salted_df = salt_skewed_column(skewed_df, "skewed_column")

  3. Monitoring Query Performance: 

    def monitor_query_performance(query):
    # Enable detailed metrics
    spark.conf.set("spark.sql.execution.metrics.enabled", "true")

    # Execute with different configurations
    results = {}
    for aqe in [True, False]:
        for dpp in [True, False]:
            spark.conf.set("spark.sql.adaptive.enabled", str(aqe))
            spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", str(dpp))

            start_time = time.time()
            spark.sql(query).collect()
            execution_time = time.time() - start_time

            results[f"AQE_{aqe}_DPP_{dpp}"] = execution_time

    return results

DPP and AQE significantly improve Spark query performance when properly implemented. Regular monitoring and tuning of these features ensure optimal performance as data volumes and query patterns evolve. The key lies in understanding your data characteristics and query patterns to leverage these optimizations effectively.

Want to receive update about our upcoming podcast?

Thanks for joining our newsletter.
Oops! Something went wrong.

Latest Articles

View All Articles

Implementing custom serialization and deserialization in Apache Kafka for optimized event processing performance

Dive deep into implementing custom serialization and deserialization in Apache Kafka to optimize event processing performance. This comprehensive guide covers building efficient binary serializers, implementing buffer pooling for reduced garbage collection, managing schema versions, and integrating compression techniques. With practical code examples and performance metrics, learn how to achieve up to 65% higher producer throughput, 45% better consumer throughput, and 60% reduction in network bandwidth usage. Perfect for developers looking to enhance their Kafka implementations with advanced serialization strategies.

time
11
 min read

Designing multi-agent systems using LangGraph for collaborative problem-solving

Learn how to build sophisticated multi-agent systems using LangGraph for collaborative problem-solving. This comprehensive guide covers the implementation of a software development team of AI agents, including task breakdown, code implementation, and review processes. Discover practical patterns for state management, agent communication, error handling, and system monitoring. With real-world examples and code implementations, you'll understand how to orchestrate multiple AI agents to tackle complex problems effectively. Perfect for developers looking to create robust, production-grade multi-agent systems that can handle iterative development workflows and maintain reliable state management.

time
7
 min read

Designing event-driven microservices architectures using Apache Kafka and Kafka Streams

Dive into the world of event-driven microservices architecture with Apache Kafka and Kafka Streams. This comprehensive guide explores core concepts, implementation patterns, and best practices for building scalable distributed systems. Learn how to design event schemas, process streams effectively, and handle failures gracefully. With practical Java code examples and real-world architectural patterns, discover how companies like Netflix and LinkedIn process billions of events daily. Whether you're new to event-driven architecture or looking to optimize your existing system, this guide provides valuable insights into building robust, loosely coupled microservices.

time
12
 min read