🪣 Bucketing Strategies in PySpark

Bucketing Architecture

Architecture Diagram

Detailed Explanation

What is Bucketing?

Bucketing distributes data into a fixed number of files (buckets) based on hash value of columns.

Key Takeaway: Unlike partitioning (directory structures), bucketing organizes data within a directory into hash-based file groups.

Core Benefits

Hash Function and Bucket Assignment

• Uses murmur3 hash (Spark's default) producing 32-bit integer
• Bucket index: hash(column_value) % num_buckets
• All rows with same bucketing column value → same bucket file

Bucketed Join Logic:

1. Both tables must have same number of buckets
2. Each bucket in Table A joins with corresponding bucket in Table B
3. No shuffle required—bucket-aware sort-merge join

Bucket Pruning Optimization

When query includes equality predicate on bucketing column:

Performance Impact: Can skip 99%+ of data for point queries.

Number of Buckets

Calculation: num_buckets = total_data_size / target_file_size

Important: Bucket count must match between tables for bucketed joins.

Sorted Bucketing

Additional optimization within each bucket:

• Data sorted by bucketing columns
• Enables efficient range scans
• Sort order maintained at file level
• Concatenation of files preserves sort order

Bucketing + Partitioning Interaction

Pruning: Bucket pruning only works within a partition—queries filtering on both benefit from both strategies.

Warning: Over-partitioning + bucketing can create too many small files.

Mathematical Foundations

Definition: Hash Bucketing

A hash bucketing strategy partitions dataset $D$ into $B$ buckets via a hash function $h: \text{key} \rightarrow \{0, 1, \ldots, B-1\}$. Bucket $b$ contains:

Bucket Pruning

For equi-join on bucketed columns with $B$ buckets, only matching buckets need scanning:

Join scans reduce from $|D_1| \times |D_2|$ to $B \times \frac{|D_1|}{B} \times \frac{|D_2|}{B} = \frac{|D_1| \times |D_2|}{B}$.

Load Balance Theorem

For hash function $h$ mapping $n$ keys to $B$ buckets, the expected maximum bucket size is:

Load imbalance approaches 0 as $n/B \rightarrow \infty$ for good hash functions.

Sort-Merge Join Cost

Without bucketing, sort-merge join cost is:

With bucketing, sorting is eliminated: $C_{\text{bucketed}} = O(|D_1| + |D_2|)$.

Bucket Count Selection

Optimal bucket count $B$ minimizes total I/O:

where $S_{\text{block}}$ is the HDFS block size.

Key Insight

Bucketing trades write-time computation for read-time savings. The benefit is maximal when the same bucketed column is used for joins across multiple queries. Over-bucketing ($B > |D|/S_{\text{block}}$) creates too many small files, hurting I/O.

Summary

Hash bucketing partitions data by key, enabling bucket pruning that reduces join scan by factor $B$. Load balance follows a balls-into-bins model. Optimal bucket count balances block-aligned I/O against overhead. Bucketing is most beneficial for repeated equi-joins on the same key.

Key Concepts Table

Code Examples

Creating Bucketed Tables

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

spark = SparkSession.builder \
    .appName("BucketingStrategies") \
    .config("spark.sql.sources.bucketing.enabled", "true") \
    .config("spark.sql.sources.bucketing.maxBuckets", "100000") \
    .getOrCreate()

# Create large dataset for bucketing demonstration
from pyspark.sql.functions import rand, floor, concat, lit

# Generate customer orders (10M rows)
orders_df = spark.range(0, 10_000_000) \
    .withColumn("order_id", concat(lit("ORD-"), col("id"))) \
    .withColumn("customer_id", concat(lit("CUST-"), (col("id") % 1_000_000))) \
    .withColumn("product_id", concat(lit("PROD-"), (col("id") % 50_000))) \
    .withColumn("amount", (rand() * 1000 + 10).cast("decimal(10,2)")) \
    .withColumn("order_date", date_add(lit("2024-01-01"), (col("id") % 365).cast("int")))

# Create customers dataset (1M rows)
customers_df = spark.range(0, 1_000_000) \
    .withColumn("customer_id", concat(lit("CUST-"), col("id"))) \
    .withColumn("name", concat(lit("Customer_"), col("id"))) \
    .withColumn("segment", 
        when(rand() < 0.3, "Premium")
        .when(rand() < 0.6, "Standard")
        .otherwise("Basic"))

# Write bucketed tables
# Bucket by customer_id with 256 buckets
orders_df.write \
    .bucketBy(256, "customer_id") \
    .sortBy("customer_id") \
    .mode("overwrite") \
    .saveAsTable("bucketed_orders")

customers_df.write \
    .bucketBy(256, "customer_id") \
    .sortBy("customer_id") \
    .mode("overwrite") \
    .saveAsTable("bucketed_customers")

# Verify bucketing
spark.sql("DESCRIBE EXTENDED bucketed_orders").show(truncate=False)
spark.sql("DESCRIBE EXTENDED bucketed_customers").show(truncate=False)

Bucket Pruning Demonstration

# Enable bucket pruning for queries
spark.conf.set("spark.sql.sources.bucketing.enabled", "true")

# Query with bucket pruning (equality predicate on bucket column)
# This should only scan 1/256 of the data
import time

# Point query - bucket pruning enabled
start = time.time()
result1 = spark.sql("""
    SELECT o.order_id, o.amount, c.name
    FROM bucketed_orders o
    JOIN bucketed_customers c ON o.customer_id = c.customer_id
    WHERE o.customer_id = 'CUST-000042'
""")
result1.show(truncate=False)
pruned_time = time.time() - start
print(f"Bucket pruning query time: {pruned_time:.2f} seconds")

# Query without bucket pruning (range predicate)
start = time.time()
result2 = spark.sql("""
    SELECT o.order_id, o.amount, c.name
    FROM bucketed_orders o
    JOIN bucketed_customers c ON o.customer_id = c.customer_id
    WHERE o.amount > 500
""")
result2.show(5, truncate=False)
full_scan_time = time.time() - start
print(f"Full scan query time: {full_scan_time:.2f} seconds")

# Verify bucket pruning in query plan
spark.sql("""
    EXPLAIN SELECT * FROM bucketed_orders 
    WHERE customer_id = 'CUST-000042'
""").show(truncate=False)

Bucket-Aware Joins

# Bucketed join - should NOT shuffle
start = time.time()
bucketed_join_df = spark.sql("""
    SELECT 
        o.order_id,
        o.amount,
        o.order_date,
        c.name,
        c.segment
    FROM bucketed_orders o
    JOIN bucketed_customers c 
    ON o.customer_id = c.customer_id
""")
bucketed_join_df.count()  # Force execution
bucketed_join_time = time.time() - start
print(f"Bucketed join time: {bucketed_join_time:.2f} seconds")

# Non-bucketed join for comparison
orders_df.write.mode("overwrite").saveAsTable("regular_orders")
customers_df.write.mode("overwrite").saveAsTable("regular_customers")

start = time.time()
regular_join_df = spark.sql("""
    SELECT 
        o.order_id,
        o.amount,
        o.order_date,
        c.name,
        c.segment
    FROM regular_orders o
    JOIN regular_customers c 
    ON o.customer_id = c.customer_id
""")
regular_join_df.count()  # Force execution
regular_join_time = time.time() - start
print(f"Regular join time: {regular_join_time:.2f} seconds")

print(f"Speedup: {regular_join_time / bucketed_join_time:.1f}x")

# Verify no shuffle in bucketed join
spark.sql("""
    EXPLAIN SELECT * FROM bucketed_orders o
    JOIN bucketed_customers c ON o.customer_id = c.customer_id
""").show(truncate=False)

Multi-Column Bucketing

# Bucket by multiple columns for complex query patterns
complex_orders_df = orders_df \
    .withColumn("region", 
        when(rand() < 0.25, "North")
        .when(rand() < 0.5, "South")
        .when(rand() < 0.75, "East")
        .otherwise("West"))

# Write with multi-column bucketing
complex_orders_df.write \
    .bucketBy(16, "region", "customer_id") \
    .sortBy("region", "customer_id") \
    .mode("overwrite") \
    .saveAsTable("multi_bucket_orders")

# Query benefits from both columns
spark.sql("""
    SELECT region, customer_id, SUM(amount) as total
    FROM multi_bucket_orders
    WHERE region = 'North' AND customer_id = 'CUST-000042'
    GROUP BY region, customer_id
""").show(truncate=False)

# Verify multi-column bucketing
spark.sql("DESCRIBE EXTENDED multi_bucket_orders").show(truncate=False)

Performance Metrics

Best Practices

1. Match bucket counts between join tables to enable bucket-aware joins—mismatched counts force a shuffle even with bucketing
2. Target 128-256 MB file size per bucket by calculating num_buckets = total_data_size / target_file_size
3. Use sortBy within buckets for columns commonly used in range queries to enable efficient intra-bucket scans
4. Avoid bucketing on high-cardinality columns (e.g., UUIDs) unless you need bucket pruning—hash distribution will be uniform regardless
5. Combine partitioning and bucketing when queries commonly filter on both temporal and entity columns
6. Monitor small file counts—if buckets contain files < 10 MB, reduce the number of buckets
7. Use INSERT INTO instead of write.saveAsTable for subsequent loads to maintain bucket structure
8. Enable spark.sql.sources.bucketing.enabled=true explicitly to ensure bucket pruning is active
9. Avoid bucket evolution (changing bucket count) as it requires a full table rewrite
10. Use ANALYZE TABLE after bucketing to update column statistics for the optimizer
11. Test bucket pruning with EXPLAIN to verify the query plan shows bucket pruning (scan only matching buckets)
12. Consider bucketing for CDC workloads where upserts on the bucket column benefit from targeted file updates

See also: Adaptive Query Execution (30), JSON/XML Parsing (28), Cost Optimization (39)

See Also

• Partitioning Strategies — Data partitioning approaches
• Join Optimization — Join performance optimization
• Adaptive Query Execution — AQE and dynamic optimization
• Caching Persistence — Caching strategies