我有一个名为a的表，它有一个名为id的列。我有另一个名为b的表，它有一个名为a\id的列，它引用表a的id列（外键）。我想基于此列连接这两个表，但这两个表都非常大（表a的大小约为1 tb，表b的大小约为4.2 tb）。因此，我的计划是尽可能减少无序处理以提高效率，为此，我尝试对Dataframe进行分区，以便两列中具有相同值的行位于同一分区中，如下面的示例所示。
示例表：

A: 
id  category
1   tennis
2   cricket

B:
name   a_id
Roger  1
Novak  1
Lara   2
Ricky  2

分区后我想实现的目标：

Partition 1:
A: 
id  category
1   tennis

B:
name   a_id
Roger  1
Novak  1

Partition 2:
A: 
id  category
2   cricket

B:
name   a_id
Lara   2
Ricky  2

我试图通过在加载表之后在列id上重新划分a和在列a\u id上重新划分b来实现这一点。我在一个简短的测试数据集上检查了这个，它似乎呈现了我想要的结果。
但在完整的数据集上，spark job在加入过程中进行了大量的洗牌，并被卡住了。那么，最有效的方法是什么呢？