🎯 PySpark Partitioning Strategies

Partitioning Comparison

Architecture Diagram

📚 Detailed Explanation

1. Why Partitioning Matters

Partitioning determines how data is distributed across the cluster.

Key Benefits:

• Performance: Reduces data shuffling and network I/O
• Parallelism: Enables concurrent processing across executors
• Data locality: Keeps related data together for efficient joins
• Query optimization: Enables partition pruning

2. Round-Robin Partitioning

Round-robin partitioning distributes data evenly across partitions regardless of key values.

Characteristics:

• Equal partition sizes
• No data locality
• Simple distribution
• Used by repartition(n)

When to Use:

• When data is evenly distributed and you need more partitions
• When you don't have a natural partition key
• For load balancing across executors

3. Hash Partitioning

Hash partitioning assigns records to partitions based on the hash of a key value.

Formula: partition = hash(key) % num_partitions

Characteristics:

• Same key → same partition
• Enables partition pruning
• May cause data skew
• Used by partitionBy(), repartition(n, col)

When to Use:

• When you frequently filter or join on a specific key
• When you want to co-locate related data
• For partition-level operations

4. Range Partitioning

Range partitioning divides data into ordered ranges.

Characteristics:

• Data ordered across partitions
• Enables range pruning
• Requires boundary calculation
• Used for ordered data

When to Use:

• When you frequently query ranges
• For time-series data
• For ordered analytics

5. Bucket Partitioning

Bucket partitioning is a file-based partitioning scheme that divides data into a fixed number of files based on hash of a key.

Characteristics:

• Fixed number of files per table
• Files are pre-partitioned by key
• Enables bucket joins (no shuffle)
• Persistent across writes

When to Use:

• When you frequently join two large tables
• When you want to avoid shuffle in joins
• For repeated join operations

6. Partition Count Optimization

The number of partitions should be based on:

• Data size: Target 128MB-200MB per partition
• Executor cores: 2-4 partitions per core
• Cluster size: Total partitions = executors × cores × 2-4
• Operation type: More partitions for parallel operations

Formula:

optimal_partitions = max(total_data_size / target_partition_size, 
                        num_executors * cores_per_executor * 2)

7. Data Skew and Mitigation

Data skew occurs when some partitions have significantly more data than others.

Symptoms:

• Longer task durations for skewed partitions
• Underutilization of cluster resources
• Memory pressure on executors processing skewed data

Detection:

• Monitor task duration in Spark UI
• Check partition sizes
• Use df.groupBy("partition").count()

Mitigation:

1. Salting: Add random prefix to skewed keys
2. Broadcast join: Avoid shuffle for small tables
3. AQE: Adaptive Query Execution for automatic skew handling

8. Partition Pruning

Partition pruning eliminates unnecessary partition reads based on filter predicates.

How it Works:

1. Filter predicate is analyzed
2. Matching partitions are identified
3. Only matching partitions are read
4. Other partitions are skipped entirely

Example:

-- Table partitioned by year, month
SELECT * FROM events 
WHERE year = 2024 AND month = 1

-- Only reads partitions for year=2024, month=1
-- Skips all other partitions

Performance Tip: Partition pruning can provide 10x–50x speedup by reducing the amount of data read.

🔑 Key Concepts Table

💻 Code Examples

Example 1: Basic Partitioning

from pyspark.sql import SparkSession
from pyspark.sql.functions import *

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

# Create sample data
df = spark.range(1000000).withColumn(
    "key", col("id") % 100
).withColumn(
    "value", concat(lit("value_"), col("id"))
)

print(f"Initial partitions: {df.rdd.getNumPartitions()}")

# Round-robin repartition (equal sizes)
# Parameter: numPartitions — target number of partitions
# Full shuffle: all data is redistributed
df_rr = df.repartition(20)
print(f"After repartition(20): {df_rr.rdd.getNumPartitions()}")

# Hash repartition by key
# Parameters: numPartitions, *cols — hash partition by specified columns
# Same key always goes to same partition
df_hash = df.repartition(20, "key")
print(f"After repartition(20, key): {df_hash.rdd.getNumPartitions()}")

# Coalesce to reduce partitions
# Parameter: numPartitions — must be less than current count
# No shuffle: combines existing partitions
df_coalesced = df_rr.coalesce(10)
print(f"After coalesce(10): {df_coalesced.rdd.getNumPartitions()}")

# Check partition sizes
def get_partition_sizes(rdd):
    return rdd.mapPartitions(lambda it: [sum(1 for _ in it)]).collect()

print(f"Partition sizes (rr): {get_partition_sizes(df_rr.rdd)[:5]}...")

Example 2: Hash Partitioning for Joins

# Create two large DataFrames
users = spark.range(1000000).withColumn(
    "user_id", col("id")
).withColumn(
    "name", concat(lit("user_"), col("id"))
)

orders = spark.range(5000000).withColumn(
    "order_id", col("id")
).withColumn(
    "user_id", col("id") % 1000000
).withColumn(
    "amount", (col("id") * 10.0) % 1000
)

# Method 1: Repartition both by join key
# Parameters: numPartitions, col — hash partition by join key
users_repartitioned = users.repartition(100, "user_id")
orders_repartitioned = orders.repartition(100, "user_id")

# Join (no shuffle if both are partitioned by same key)
result = users_repartitioned.join(orders_repartitioned, "user_id")
result.explain()

# Method 2: Use bucketing for persistent partitioning
# Parameters:
#   bucketBy(50, "user_id") — 50 hash buckets by user_id
#   sortBy("user_id") — sort within each bucket
#   saveAsTable("users_bucketed") — persist as Hive table
users.write \
    .bucketBy(50, "user_id") \
    .sortBy("user_id") \
    .saveAsTable("users_bucketed")

orders.write \
    .bucketBy(50, "user_id") \
    .sortBy("user_id") \
    .saveAsTable("orders_bucketed")

# Join bucketed tables (no shuffle!)
users_bucketed = spark.table("users_bucketed")
orders_bucketed = spark.table("orders_bucketed")
result_bucketed = users_bucketed.join(orders_bucketed, "user_id")
result_bucketed.explain()

Example 3: Range Partitioning

# Create time-series data
from pyspark.sql.types import TimestampType
from datetime import datetime, timedelta

# Generate time-series data
base_date = datetime(2024, 1, 1)
time_series = spark.range(1000000).withColumn(
    "timestamp", 
    (lit(base_date.cast("long")) + col("id") * 3600).cast("timestamp")
).withColumn(
    "value", col("id") * 1.0
)

# Range partition by timestamp
# Parameter: numPartitions — target partition count
# repartitionByRange sorts data and creates ordered partitions
df_range = time_series.repartitionByRange(4, "timestamp")

# Check distribution
df_range.groupBy(
    floor(unix_timestamp("timestamp") / (365 * 24 * 3600)).alias("year")
).count().show()

Example 4: Partition Pruning

# Create partitioned table
events = spark.range(1000000).withColumn(
    "event_id", col("id")
).withColumn(
    "event_date", date_add(lit("2024-01-01"), (col("id") % 365).cast("int"))
).withColumn(
    "user_id", col("id") % 10000
)

# Write with partitioning
# Parameter: partitionBy("event_date") — creates directory structure
# Output: event_date=2024-01-01/part-00000.parquet
events.write \
    .partitionBy("event_date") \
    .mode("overwrite") \
    .parquet("events_partitioned/")

# Read and query with partition pruning
events_df = spark.read.parquet("events_partitioned/")

# This query only reads partitions for January 2024
result = events_df.filter(
    col("event_date").between("2024-01-01", "2024-01-31")
)

# Check physical plan for partition pruning
result.explain()

# Verify partition pruning in Spark UI
# Look for "PartitionFilters" in the scan

📊 Performance Metrics

✅ Best Practices

1. Choose Right Partition Count

# Target 128MB-200MB per partition
data_size_gb = 10  # 10GB
target_partition_mb = 128
optimal_partitions = int(data_size_gb * 1024 / target_partition_mb)
print(f"Optimal partitions: {optimal_partitions}")  # ~80

# Consider executor cores
num_executors = 10
cores_per_executor = 4
max_partitions = num_executors * cores_per_executor * 4
print(f"Max partitions: {max_partitions}")  # 160

2. Use Hash Partitioning for Joins

# Partition both DataFrames by join key
# Parameter: numPartitions, col — hash partition
df1_partitioned = df1.repartition(100, "join_key")
df2_partitioned = df2.repartition(100, "join_key")

# Join (no shuffle)
result = df1_partitioned.join(df2_partitioned, "join_key")

3. Use Bucketing for Repeated Joins

# Write bucketed tables
# bucketBy(numBuckets, *cols) — hash into fixed buckets
# sortBy(*cols) — sort within buckets
df1.write.bucketBy(100, "key").sortBy("key").saveAsTable("t1_bucketed")
df2.write.bucketBy(100, "key").sortBy("key").saveAsTable("t2_bucketed")

# Join bucketed tables (no shuffle)
result = spark.table("t1_bucketed").join(spark.table("t2_bucketed"), "key")

4. Partition Large Tables for Query Optimization

# Partition by frequently filtered columns
# Parameter: partitionBy(*cols) — creates directory hierarchy
df.write.partitionBy("year", "month").parquet("output/")

# Enables partition pruning
df_filtered = spark.read.parquet("output/").filter(
    (col("year") == 2024) & (col("month") == 1)
)

5. Monitor Partition Distribution

# Check partition sizes
partition_sizes = df.rdd.mapPartitions(
    lambda it: [sum(1 for _ in it)]
).collect()

print(f"Partition sizes: {partition_sizes}")
print(f"Min: {min(partition_sizes)}, Max: {max(partition_sizes)}")
print(f"Skew ratio: {max(partition_sizes) / min(partition_sizes):.2f}")

6. Avoid Over-Partitioning

# BAD: Too many small partitions
df.repartition(10000)  # Creates 10K tiny partitions

# GOOD: Right-sized partitions
df.repartition(100)  # Creates 100 manageable partitions

See Also

• Joins Optimization
• Caching and Persistence
• Partitioning Strategies
• Kafka Streams
• Streaming Pipelines