How to avoid shuffles while joining DataFrames on unique keys?

Question

I have two DataFrames A and B:

• A has columns (id, info1, info2) with about 200 Million rows
• B only has the column id with 1 million rows

The id column is unique in both DataFrames.

I want a new DataFrame which filters A to only include values from B.

if B was very small I know I would something along the lines of

A.filter($("id") isin B("id"))

but B is still pretty large, so not all of it can fit as a broadcast variable.

and I know I could use

A.join(B, Seq("id"))

but that wouldn't harness the uniqueness and I'm afraid will cause unnecessary shuffles.

What is the optimal method to achieve that task?

Answer 1

If you have not applied any partitioner on Dataframe A, May be this will help you understanding Join And Shuffle concepts.

Without Partitioner :

A.join(B, Seq("id"))

By default, this operation will hash all the keys of both dataframes, sending elements with the same key hash across the network to the same machine, and then join together the elements with the same key on that machine. Here you have to notice that both dataframes shuffle across the network.

With HashPartitioner: Call partitionBy() when building A Dataframe, Spark will now know that it is hash-partitioned, and calls to join() on it will take advantage of this information. In particular, when we call A.join(B, Seq("id")), Spark will shuffle only the B RDD. Since B has less data than A you don't need to apply partitioner on B

ex:

 val A = sc.sequenceFile[id, info1, info2]("hdfs://...")
     .partitionBy(new HashPartitioner(100)) // Create 100 partitions
     .persist()
 A.join(B, Seq("id"))

Reference is from Learning Spark book.

Answer 2

My default advice on how to optimize joins is:

1. Use a broadcast join if you can (From your question it seems your tables are large and a broadcast join is not an option). One option in Spark is to perform a broadcast join (aka map-side join in hadoop world). With broadcast join, you can very effectively join a large table (fact) with relatively small tables (dimensions) by avoiding sending all data of the large table over the network.

   You can use broadcast function to mark a dataset to be broadcast when used in a join operator. It uses spark.sql.autoBroadcastJoinThreshold setting to control the size of a table that will be broadcast to all worker nodes when performing a join.

2. Use the same partitioner. If two RDDs have the same partitioner, the join will not cause a shuffle. Note however, that the lack of a shuffle does not mean that no data will have to be moved between nodes. It's possible for two RDDs to have the same partitioner (be co-partitioned) yet have the corresponding partitions located on different nodes (not be co-located). This situation is still better than doing a shuffle, but it's something to keep in mind. Co-location can improve performance, but is hard to guarantee.

3. If the data is huge and/or your clusters cannot grow such that even (2) above leads to OOM, use a two-pass approach. First, re-partition the data and persist using partitioned tables (dataframe.write.partitionBy()). Then, join sub-partitions serially in a loop, "appending" to the same final result table.

Answer 3

If I understand your question correctly, you want to use a broadcast join that replicates DataFrame B on every node so that the semi-join computation (i.e., using a join to filter id from DataFrame A) can compute independently on every node instead of having to communicate information back-and-forth between each other (i.e., shuffle join).

You can run join functions that explicitly call for a broadcast join to achieve what you're trying to do:

import org.apache.spark.sql.functions.broadcast

val joinExpr = A.col("id") === B.col("id")

val filtered_A = A.join(broadcast(B), joinExpr, "left_semi")

You can run filtered_A.explain() to verify that a broadcast join is being used.