PySpark 优化两个大表的左连接

Question

I'm using the most updated version of PySpark on Databricks. I have two tables each of the size ~25-30GB. I want to join Table1 and Table2 at the "id" and "id_key" columns respectively. I'm able to do that with the command below but when I run my spark job the join is skewed resulting in +95% of my data on one executor making the job take forever. This is happening when I attempt to load the data after transforming it.

Table1 has 13 total columns where the "id" column has a lot of null values with some actual id values.

Table2 has 3 total columns where "id_key" has all possible id values each appearing once.

I tried broadcasting but because the tables are pretty large, I get OutOfMemory errors when running the job

Table1.join(Table2, Table1.id == Table2.id_key, "left")

I'm thinking of salting but not sure how to go about that or if it is the right solution.

Answer 1

As I understood your problem, I guess Spark must put all rows with null id in Table1 in the same partition during join's partitioning.

To solve this, you can use the following pattern:

• split Table1 dataframe between null ids dataframe and not null ids dataframe
• join your not null ids dataframe with Table2
• add columns of Table2 with null value to null ids dataframe
• union the obtained not null ids dataframe and null ids dataframe

You can find below the code translation of this pattern:

from pyspark.sql import functions as F

Table1_with_null_ids = Table1.filter(F.col('id').isNull())
Table1_with_not_null_ids = Table1.filter(F.col('id').isNotNull())

Table1_with_not_null_ids_joined = Table1_with_not_null_ids.join(
  Table2, 
  Table1_with_not_null_ids.id == Table2.id_key, 
  'left'
)

Table1_with_null_ids_joined = Table1_with_null_ids \
  .withColumn('id_key', F.lit(None)) \
  .withColumn('table2_column2', F.lit(None)) \
  .withColumn('table2_column3', F.lit(None))

Table1_joined = Table1_with_not_null_ids_joined.unionByName(Table1_with_null_ids_joined)

It avoids doing manual salting, and it may improve performance as you join Table2 with much less rows on Table1 side

However you need to compute the input Table1 two times as you perform filter two times on the same Table1 . If computing input Table1 is an expensive process, you can either cache Table1 before double filtering or proceed as you suggested and add a salting column to Table1 and Table2 and use it in your join expression.