[SPARK-4433] fix a racing condition in zipWithIndex

[mengxr]

Spark hangs with the following code:

sc.parallelize(1 to 10).zipWithIndex.repartition(10).count()

This is because ZippedWithIndexRDD triggers a job in getPartitions and it causes a deadlock in DAGScheduler.getPreferredLocs (synced). The fix is to compute startIndices during construction.

This should be applied to branch-1.0, branch-1.1, and branch-1.2.

@pwendell

[SparkQA]

Test build #23433 has started for PR 3291 at commit c284d9f.

• This patch merges cleanly.

[SparkQA]

Test build #23433 has finished for PR 3291 at commit c284d9f.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[AmplabJenkins]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins//job/SparkPullRequestBuilder/23433/
Test PASSed.

[rxin]

This makes the whole thing eager, isn't it?

[mengxr]

One way to do this lazily is via shuffle:

      val identityPartitioner = new Partitioner {
        override def numPartitions: Int = p
        override def getPartition(key: Any): Int = key.asInstanceOf[Int]
      }
      val startIndices = PartitionPruningRDD.create(this, _ < p - 1) // skip the last partition
        .mapPartitionsWithIndex { (split, iter) =>
          val size = Utils.getIteratorSize(iter)
          Iterator.range(split + 1, p).map { i =>
            (i, size)
          }
        }.reduceByKey(identityPartitioner, _ + _)
        .values
      this.zipPartitions(startIndices) { (iter, startIndexIter) =>
        val startIndex = if (startIndexIter.hasNext) startIndexIter.next() else 0L
        iter.zipWithIndex.map { case (item, localIndex) =>
          (item, startIndex + localIndex)
        }
      }

But I think this is more expensive.

[mengxr]

@rxin Does it look good to you? I hope that this fix can get into 1.1.1.

[rxin]

lgtm

[mengxr]

Thanks! Merged into master, branch-1.2, 1.1, and 1.0.

[andrewor14]

Hey @mengxr can you close this