We shouldn't cache data in our internal methods without giving users control over it. For this method, I'd suggest the following, simpler/faster implementation:
- Pass 1: use `self.mapPartitionsWithIndex((index, iter) => (index, iter.size)).collect()` to get the length of each partition. From this, you can get a (partitionId, numElementsBefore) table. Broadcast this as an Array[Long] or Map[Int, Long]. You can do this imperatively with just a loop, no need for map and reduce there.
- Pass 2: use self.mapPartitionsWithIndex again to take each partition, add the numElementsBefore from this table, and return record indices (same as below)

Right now you're doing some of pass 1 in parallel with this map/reduce and range stuff, and it's not clear that it's O(n) in computing the numElementsBefore, but it could all just be done locally on the driver program.