Broadcast Joins, Shared Variables & Bucketing

In an interview you may see a query plan with an Exchange line. An Exchange is a shuffle: Spark moves rows between machines. The senior question is: how would you remove this shuffle?

You already know two joins from "Joins in Spark": the sort-merge join and the broadcast hash join. This chapter adds the tools that chapter did not cover.

1. Broadcast variables: share a small Python value with every task.
2. Accumulators: count things inside tasks, and read the total on the driver.
3. Bucketing: store a table so that later joins need no shuffle.
4. A table of ways to remove a shuffle, from the cheapest to the most work.

We use the shop again: 6 orders and 5 customers. Customer 5, Omar, has no orders. Customer 4, Dev, has a NULL country.

"Joins in Spark" taught the broadcast hash join. Spark copies a small table to every executor, and the big table does not move. In the plan you see BroadcastExchange and no Exchange hashpartitioning line.

Keep two limits in mind.

• The driver collects the small table first, so the driver needs memory for it.
• Every executor keeps a full copy, so a table that is too big can run out of memory.

A broadcast join is not a broadcast variable. The names sound alike, but you use them in different places.

Sometimes your own Python function needs a lookup table. A UDF (user-defined function) is a Python function that you use on a column. rdd.map() runs your function on each item of an RDD.

Both run inside tasks. If the function uses a dict, Spark sends a copy of the dict with the function to every task. A broadcast variable sends it once to each executor, and tasks on that executor reuse it.

Here the lookup turns a country code into a country name. Look at the code and the result.

from pyspark.sql import functions as F

sc = spark.sparkContext
names = sc.broadcast({"IN": "India", "GB": "United Kingdom", "US": "United States"})

@F.udf("string")
def country_name(code):
    return names.value.get(code)

customers.select("name", country_name("country").alias("country_name")).show()

Treat the value as read-only. A change inside a task is never sent back to the driver. Other tasks on that executor may see it, so results become unsafe.

If the lookup is a small table, a broadcast join is often better. Spark can plan a join. It cannot look inside a UDF.

Memory is the main limit. Clean up when you are done.

• The driver holds the value, and every executor holds a copy. In PySpark, each Python process that runs tasks loads its own copy too.
• A big value can run the driver or an executor out of memory.
• Call names.unpersist() to remove the copies from the executors. Spark sends them again if a task needs them.
• Call names.destroy() to remove the variable everywhere. You cannot use it after that.

An accumulator sends data the other way: from tasks to the driver. Tasks add to a counter, and you read the total on the driver. Tasks cannot read the total.

Use it to count things you want to watch, such as orders with no amount. Look at the code and the result.

missing = spark.sparkContext.accumulator(0)

def check(row):
    if row.amount is None:
        missing.add(1)

orders.foreach(check)          # foreach is an action: it runs the tasks
print("orders with no amount:", missing.value)

Spark waits for an action before it runs any task.

There is one big risk. A task can run again: after a failure, or when a second action recomputes the same DataFrame. Then its adds can count twice.

So use accumulators for debugging and monitoring. Do not use one to build a result that you store.

Bucketing splits a table into a fixed number of buckets when you write it. Rows with the same key go to the same bucket. Spark can then join two such tables without a shuffle.

Spark keeps the bucket layout in the catalog, which is the list of tables that Spark knows. So you need saveAsTable. A plain path has no place to keep the layout.

We use 3 buckets and the model from the joins chapter: bucket = customer_id % 3. Real Spark uses a hash, so real buckets differ. Here is the placement.

Here is the code that writes both tables in this layout.

(orders.write.bucketBy(3, "customer_id").sortBy("customer_id")
       .saveAsTable("orders_bucketed"))

(customers.write.bucketBy(3, "customer_id").sortBy("customer_id")
          .saveAsTable("customers_bucketed"))

"Joins in Spark" showed the plain plan: two Exchange lines, one for each table.

Now join the two bucketed tables. The first line turns off broadcast, because our tables are tiny.

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)
o = spark.table("orders_bucketed")
c = spark.table("customers_bucketed")
o.join(c, "customer_id").explain()

Here is the plan. It is simplified. Real output also has AdaptiveSparkPlan and Project lines, Filter isnotnull lines and column ids.

== Physical Plan (simplified) ==
SortMergeJoin [customer_id], [customer_id], Inner
:- Sort [customer_id ASC]
:  +- FileScan parquet orders_bucketed (Bucketed: true)
+- Sort [customer_id ASC]
   +- FileScan parquet customers_bucketed (Bucketed: true)

The join is still a SortMergeJoin. It only skips the move.

The Sort lines stay. In Spark 3, a bucketed scan does not tell Spark that the rows are sorted. Bucketing saves the shuffle, not the sort.

Bucketing also helps a groupBy on the same key, because equal keys already sit in one bucket. You pay once, when you write the table.

It helps for a big table that you join many times on the same key. For a join you run once, it is not worth the write.

Bucketing only works when the layouts fit. Check these six points.

1. The key must match. Tables bucketed on customer_id help joins on customer_id only.
2. The bucket counts must match. Tables with 3 and 4 buckets cannot join without a shuffle.
3. If only one table is bucketed, both tables usually shuffle. Only the other table shuffles when the bucket count is at least spark.sql.shuffle.partitions (200 by default).
4. Hive, an older SQL engine, hashes buckets differently from Spark. Bucket both tables with Spark.
5. The setting spark.sql.sources.bucketing.enabled must be on (it is by default). Otherwise Spark ignores the buckets.
6. Delta Lake tables do not support bucketBy. The chapter Delta Lake & the Lakehouse covers Delta Lake.

Newer Spark can merge buckets when the counts differ, such as 8 and 4. The setting is spark.sql.bucketing.coalesceBucketsInJoin.enabled.

It needs Spark 3.1 or later, and it is off by default. The larger count must be a multiple of the smaller one. It must also be at most 4 times the smaller one.

Another risk is many small files. Each write task makes one file for every bucket it holds. Look at the 3 write tasks for orders.

That is 5 files for 3 buckets. Buckets 0 and 1 have two files each. Bucket 2 has one.

After a shuffle there are 200 partitions by default. 200 write tasks × 64 buckets can make up to 12,800 files.

The fix is to repartition by the key first, with the bucket count. Then each task holds the rows of one bucket.

(orders.repartition(3, "customer_id")
       .write.mode("overwrite").bucketBy(3, "customer_id").sortBy("customer_id")
       .saveAsTable("orders_bucketed"))

"Partitioning & Parallelism" showed partitionBy("order_date"). It writes one folder for each date: 6 folders here, with one order in each. A filter on the date can skip folders.

Bucketing is different. It has a fixed number of buckets, not one folder for each value. Compare the two.

You can use both on one table. Each date folder then holds its own bucket files. Six dates and 3 buckets give up to 18 files for each write task.

Back to the senior question. First find the Exchange line in the plan. Then try the cheapest idea first.

One idea is new here: pre-aggregate. Add up the orders for each customer before the join. The join then moves fewer rows.

from pyspark.sql import functions as F

totals = orders.groupBy("customer_id").agg(F.sum("amount").alias("total"))
totals.join(customers, "customer_id").select("name", "total").show()

The groupBy still shuffles. But Spark adds up inside each partition first, so fewer rows move. The result is also already grouped by customer_id, so Spark can often reuse that layout for the join.

Here are five ideas. The first four go from the cheapest to the most work.

Four points to keep.

• A broadcast variable shares a Python value. A broadcast join shares a table.
• Accumulators count inside tasks. Trust them for debugging, not for stored results.
• Bucketing moves the work to the write. Both tables need the same key and bucket count.
• To remove a shuffle, go from cheap to costly: filter, broadcast, pre-aggregate, bucket.

Answer these without looking back.

1. What is the difference between a broadcast join and a broadcast variable?
2. Why can an accumulator show a wrong count inside a transformation?
3. Why does bucketBy need saveAsTable and not a plain path?
4. Two tables have 3 and 4 buckets on customer_id. Can they join without a shuffle? Why?