Is the partition the partition number in the RDD, or the task number in the stages? Some stages run on just a subset of an RDD. This should be documented to prevent confusion.