Spark TaskContext两种分区ID获取方法区别及0-199分区ID获取方案

能否获取0~199的连续分区ID #

可以，你对RDD调用repartition(200)之后，生成的新RDD的分区ID就是连续的0到199的整数，完全符合你的取值要求。

两种分区ID获取方法的区别 #

两种方法本质功能完全一致，仅写法和版本兼容性有差异：

• TaskContext.get().partitionId()：全Spark版本兼容的写法，先通过静态get()方法获取当前任务的上下文实例，再调用实例方法拿到分区ID，所有支持Java API的Spark版本都可以正常运行。
• TaskContext.partitionId()：Spark 2.3及以上版本新增的静态快捷方法，内部封装了get().partitionId()的逻辑，写法更简洁，但在2.3以下版本的Spark环境中运行会抛出方法不存在的异常。

适配业务场景的代码示例 #

import org.apache.spark.TaskContext;
import scala.Tuple2;

final JavaPairRDD<Integer, Integer> baseEntity = ....repartition(200);
baseEntity.mapPartitionsToPair(entities -> {
   // 全版本兼容写法，Spark2.3+环境可直接替换为简化写法
   int partitionId = TaskContext.get().partitionId();
   return Iterators.transform(entities, entity ->{
       // 此处可直接使用拿到的partitionId，同个分区内的所有元素对应同一个ID
       return new Tuple2<>(entity._1, entity._2);
   });
});

注意事项 #

分区ID获取的逻辑只能写在算子内部的执行代码中，不能在Driver端的算子外部调用相关方法，否则会因为不存在TaskContext实例抛出运行时异常。

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