Spark Tuning - Advanced Partitioning: Hands-on Bucketing Workbook

This workbook walks you through writing and validating bucketed tables to avoid shuffles.

1. Bucketing Operations

• Task: Write a PySpark DataFrame to a hive catalog as bucketed tables and inspect the physical join plan.

2. Tasks

Task 1: Write Bucketed Tables

Write the PySpark code to bucket two DataFrames (users and orders) by user_id into 16 buckets, write them to disk, join them, and verify that the join execution runs without a shuffle.

Task 2: Implement Salting Transformations

Assume users table has a highly skewed key "GUEST_USER". Write the PySpark code to append random salts (0 to 3) to "GUEST_USER" records and join it safely with the salted orders table.

3. Step-by-Step Solutions

Solution 1: Bucketed Join Implementation

• PySpark Code:

# 1. Write users table bucketed
users_df.write.format("parquet") \
        .bucketBy(16, "user_id") \
        .sortBy("user_id") \
        .saveAsTable("bucketed_users")

# 2. Write orders table bucketed
orders_df.write.format("parquet") \
        .bucketBy(16, "user_id") \
        .sortBy("user_id") \
        .saveAsTable("bucketed_orders")

# 3. Read tables back from Catalog
b_users = spark.read.table("bucketed_users")
b_orders = spark.read.table("bucketed_orders")

# 4. Join bucketed tables (must share exact number of buckets and sorting key)
joined_df = b_users.join(b_orders, "user_id")

# 5. Explain Physical Plan to confirm absence of 'Exchange' (Shuffle)
joined_df.explain()

• Physical Plan Verification: The resulting plan shows a SortMergeJoin with direct file scans, and no Exchange (shuffle) operators in the execution path, proving bucketing worked perfectly!

Solution 2: Salting Code

import pyspark.sql.functions as F

# Add random salt (0 to 3) to skewed keys
salted_users = users_df.withColumn(
    "salted_key", 
    F.when(F.col("user_id") == "GUEST_USER", 
           F.concat(F.col("user_id"), F.lit("_"), F.randint(0, 3)))
     .otherwise(F.col("user_id"))
)

# Explode lookup order records
exploded_orders = orders_df.withColumn("salt_array", F.array([F.lit(i) for i in range(4)])) \
                            .withColumn("salt_val", F.explode("salt_array")) \
                            .withColumn("salted_order_key", 
                                        F.when(F.col("user_id") == "GUEST_USER", 
                                               F.concat(F.col("user_id"), F.lit("_"), F.col("salt_val")))
                                         .otherwise(F.col("user_id")))

# Join on salted keys
joined_df = salted_users.join(exploded_orders, salted_users.salted_key == exploded_orders.salted_order_key)