Spark Join优化：如何避免多次Shuffle实现多左连接

问题背景 #

现有两个Spark DataFrame：

• tickets：包含account_id、assignee_id、requester_id三列
• users：包含account_id、id、name三列

需要关联生成包含account_id、assignee_id、assignee_name、requester_id、requester_name的结果DataFrame。初始实现采用两次左外连接，但会触发tickets和users的两次Shuffle操作，现需解决：

1. 是否可通过Spark DataFrame API，无需复杂自定义逻辑避免双重Shuffle？
2. 验证提出的自定义分区思路是否可行

初始实现代码 #

初始方案通过两次left_outer join实现，代码如下：

tickets
    .as("ticket")
    .join(
        users.as("assignee"),
        $"ticket.account_id" === $"assignee.account_id" && $"ticket.assignee_id" === $"assignee.id",
        "left_outer"
    )
    .join(
        users.as("requester"),
        $"ticket.account_id" === $"requester.account_id" && $"ticket.requester_id" === $"requester.id",
        "left_outer"
    )
    .select(
        $"ticket.*",
        $"assignee.name".as("assignee_name"),
        $"requester.name".as("requester_name")
    )

该方案的问题在于：两次独立的join会分别触发Shuffle，导致相同数据被重复分发，浪费集群资源。

无需自定义逻辑的优化方案 #

如果users数据集规模较小，可以采用**广播连接（Broadcast Join）**来避免Shuffle：

1. 将users数据集广播到所有Executor节点，这样tickets无需进行Shuffle，两次join都可以在本地完成。
2. 实现时只需在join前对users调用broadcast方法，Spark会自动优化为广播连接：

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

tickets
    .as("ticket")
    .join(
        broadcast(users).as("assignee"),
        $"ticket.account_id" === $"assignee.account_id" && $"ticket.assignee_id" === $"assignee.id",
        "left_outer"
    )
    .join(
        broadcast(users).as("requester"),
        $"ticket.account_id" === $"requester.account_id" && $"ticket.requester_id" === $"requester.id",
        "left_outer"
    )
    .select(
        $"ticket.*",
        $"assignee.name".as("assignee_name"),
        $"requester.name".as("requester_name")
    )

该方案无需自定义分区或复杂逻辑，Spark会自动复用广播的users数据，避免两次Shuffle。若users数据量较大（超过默认广播阈值约10MB），则不适合此方案。

自定义分区思路的可行性验证 #

思路核心 #

通过自定义分区策略，让tickets与对应的assignee、requester用户数据落在同一个Executor分区内，从而在本地完成两次join，避免全局Shuffle：

1. 将总分区数p拆分为sqrt(p) × sqrt(p)的子分区网格
2. 对tickets：根据assignee_id和requester_id的哈希值确定其所在的网格分区
3. 对users：复制用户数据并分发到所有可能需要它的分区（即所有以该用户为assignee或requester的ticket所在的分区）

分区逻辑代码 #

Ticket自定义分区函数

def partitionTicket(ticket: Row, p: Int): Int = {
    val maxSubP = Math.floor(Math.sqrt(p)).toInt
    val assigneeId = ticket.getAs[Long]("assignee_id")
    val requesterId = ticket.getAs[Long]("requester_id")
    // 哈希函数可根据实际数据类型调整，这里用简单的取模示例
    val assigneeP = Math.abs(assigneeId.hashCode()) % maxSubP
    val requesterP = Math.abs(requesterId.hashCode()) % maxSubP
    
    assigneeP * maxSubP + requesterP
}

User自定义分区函数

// 先对users数据进行复制，添加target_p列（范围0到2*sqrt(p)-1）
def partitionUser(user: Row, p: Int): Int = {
    val maxSubP = Math.floor(Math.sqrt(p)).toInt
    val userId = user.getAs[Long]("id")
    val targetP = user.getAs[Int]("target_p")
    val assigneeSectionIndex = Math.abs(userId.hashCode()) % maxSubP
     
    if (targetP >= assigneeSectionIndex && targetP < assigneeSectionIndex + maxSubP) {
        assigneeSectionIndex * maxSubP + targetP % maxSubP
    } else {
        val correctedTargetP = if (targetP < assigneeSectionIndex) {
            targetP
        } else {
            targetP - maxSubP + 1
        }
        correctedTargetP * maxSubP + assigneeSectionIndex
    }
}

可行性分析 #

1. 逻辑可行性：该策略确实能让ticket与对应的assignee、requester用户落在同一分区，从而可以在本地完成两次left join，避免全局Shuffle。
2. 开销权衡：用户数据需要复制2*sqrt(p)份，当分区数p较大时（如p=100，复制20倍），会显著增加存储和内存开销，需要结合集群资源和数据规模权衡是否采用。
3. 实现复杂度：需要将DataFrame转换为RDD来应用自定义分区器，再转回DataFrame，虽然可行，但比广播连接的实现复杂度高。
4. 数据倾斜风险：哈希函数的均匀性直接影响分区数据分布，若哈希值分布不均，会导致部分分区数据量过大，引发性能瓶颈。

内容的提问来源于stack exchange，提问作者Florent