Spark Java API实现DataFrame同ID行单独分区的方法

方案一：确保每个分区仅包含单一ID的行 #

默认repartition("the_id_value")基于哈希分区，会把哈希值冲突的不同ID放入同一分区。要实现单ID独占一个分区，可按以下步骤操作：

1. 先提取所有唯一ID并收集到本地集合
2. 以唯一ID的数量作为分区数，结合ID列重分区

Java代码示例：

// 提取所有唯一ID
List<String> uniqueIds = df.select("the_id_value")
                          .distinct()
                          .as(Encoders.STRING())
                          .collectAsList();

// 按唯一ID数量重分区，保证每个ID独占一个分区
Dataset<Row> singleIdPerPartitionDf = df.repartition(uniqueIds.size(), functions.col("the_id_value"));

注意：该方案仅适合唯一ID数量较少的场景（比如几千级以内），若ID数量过大，会导致分区过多，引发Task调度压力或内存问题。

方案二：分区内按ID分组处理（适配大数量ID场景） #

如果唯一ID数量极大，直接创建对应分区会耗尽集群资源，可先通过哈希分区聚合相同ID的行，再在分区内按ID分组处理：

Java代码示例：

// 先按ID哈希分区（根据集群资源设置合理的分区数，比如200）
Dataset<Row> repartitionedDf = df.repartition(200, functions.col("the_id_value"));

// 遍历每个分区，在分区内按ID分组生成对象
repartitionedDf.foreachPartition(iterator -> {
    // 用Map存储每个ID对应的行数据
    Map<String, List<Row>> idRowMap = new HashMap<>();
    while (iterator.hasNext()) {
        Row row = iterator.next();
        String id = row.getAs("the_id_value");
        idRowMap.computeIfAbsent(id, k -> new ArrayList<>()).add(row);
    }

    // 遍历每个ID分组，生成目标对象
    for (Map.Entry<String, List<Row>> entry : idRowMap.entrySet()) {
        String targetId = entry.getKey();
        List<Row> idRows = entry.getValue();
        // 替换为你的对象转换逻辑
        YourTargetObject obj = convertRowsToObject(targetId, idRows);
        // 替换为你的对象处理逻辑
        handleTargetObject(obj);
    }
});

优势：既保证了相同ID的行在同一分区，又避免了分区数量爆炸，同时通过内存分组实现了单ID对应对象的生成。

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