This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 64badc44c272 [SPARK-57342][CORE] Reuse start indices from an ancestor
ZippedWithIndexRDD
64badc44c272 is described below
commit 64badc44c272eb644da132fba36544ea88898813
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Mon Jun 22 09:59:31 2026 -0700
[SPARK-57342][CORE] Reuse start indices from an ancestor ZippedWithIndexRDD
### What changes were proposed in this pull request?
This PR lets `ZippedWithIndexRDD` reuse partition start indices that an
ancestor has already computed, instead of always launching a counting job.
`getAncestorWithSamePartitionSizes` now stops at a `ZippedWithIndexRDD`
ancestor whose `startIndices` are already populated. The field is `transient`,
so the `startIndices != null` guard lets a deserialized ancestor fall through
to a counting job.
When the resolved ancestor is itself a `ZippedWithIndexRDD`, `startIndices`
reuses that ancestor's indices directly and skips the counting job entirely
(re-checking `!= null` for safety).
### Why are the changes needed?
Computing start indices currently runs a job to count every partition's
size. When an ancestor with identical partition sizes has already done this
work, that job is redundant; reusing the ancestor's result avoids the extra
pass over the data.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added two `RDDSuite` tests that use a job-counting `SparkListener`:
- `zipWithIndex reuses start indices from an ancestor ZippedWithIndexRDD` —
chaining `zipWithIndex` directly, and through a size-preserving `map`, reuses
the ancestor's start indices and submits no extra job, with correct indices.
- `zipWithIndex recomputes start indices when an intervening op changes
partition sizes` — a non-size-preserving op (`filter`) still triggers a
counting job and yields correct indices.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (model: claude-opus-4-8)
Closes #56403 from zhengruifeng/core_zip_zipwithindex_reuse3-dev3.
Lead-authored-by: Ruifeng Zheng <[email protected]>
Co-authored-by: Wenchen Fan <[email protected]>
Co-authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../org/apache/spark/rdd/ZippedWithIndexRDD.scala | 24 ++++++--
.../test/scala/org/apache/spark/rdd/RDDSuite.scala | 64 ++++++++++++++++++++++
2 files changed, 83 insertions(+), 5 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala
b/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala
index 60638282668d..311937952f74 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala
@@ -40,8 +40,14 @@ class ZippedWithIndexRDDPartition(val prev: Partition, val
startIndex: Long)
private[spark]
class ZippedWithIndexRDD[T: ClassTag](prev: RDD[T]) extends RDD[(T,
Long)](prev) {
+ @scala.annotation.tailrec
private def getAncestorWithSamePartitionSizes(rdd: RDD[_]): RDD[_] = {
rdd match {
+ // A ZippedWithIndexRDD ancestor that has already computed its
startIndices covers the same
+ // partition sizes, so stop here and let the caller reuse those indices
directly. startIndices
+ // is @transient and may be null on a deserialized ancestor; guard
against that so such an
+ // ancestor falls through to running a counting job instead.
+ case z: ZippedWithIndexRDD[_] if z.startIndices != null => z
case c: RDD[_] if c.isCheckpointed || c.getStorageLevel !=
StorageLevel.NONE => c
case m: MapPartitionsRDD[_, _] if m.preservesPartitionSizes =>
getAncestorWithSamePartitionSizes(m.prev)
@@ -60,11 +66,19 @@ class ZippedWithIndexRDD[T: ClassTag](prev: RDD[T]) extends
RDD[(T, Long)](prev)
Array(0L)
} else {
val ancestor = getAncestorWithSamePartitionSizes(prev)
- ancestor.context.runJob(
- ancestor,
- Utils.getIteratorSize _,
- 0 until n - 1 // do not need to count the last partition
- ).scanLeft(0L)(_ + _)
+ ancestor match {
+ // getAncestorWithSamePartitionSizes only stops at a
ZippedWithIndexRDD whose startIndices
+ // are already populated and which covers the same partition sizes as
ours, so its start
+ // indices are identical; reuse them directly and skip the counting
job. The null check is
+ // redundant with that guard but kept for safety.
+ case z: ZippedWithIndexRDD[_] if z.startIndices != null =>
z.startIndices
+ case _ =>
+ ancestor.context.runJob(
+ ancestor,
+ Utils.getIteratorSize _,
+ 0 until n - 1 // do not need to count the last partition
+ ).scanLeft(0L)(_ + _)
+ }
}
}
diff --git a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
index b9774482f949..3c295aecf274 100644
--- a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
@@ -974,6 +974,70 @@ class RDDSuite extends SparkFunSuite with
SharedSparkContext with Eventually {
assert(count === 10)
}
+ test("zipWithIndex reuses start indices from an ancestor
ZippedWithIndexRDD") {
+ val numJobs = new AtomicInteger(0)
+ val listener = new SparkListener {
+ override def onJobStart(jobStart: SparkListenerJobStart): Unit = {
+ numJobs.incrementAndGet()
+ }
+ }
+ sc.addSparkListener(listener)
+ try {
+ val n = 10
+ val expected = (0 until n).map(_.toLong)
+
+ // Building the first zipWithIndex submits one counting job: the source
has
+ // multiple partitions and no materialized ancestor to reuse.
+ val base = sc.parallelize(0 until n, 3).zipWithIndex()
+ sc.listenerBus.waitUntilEmpty()
+ assert(numJobs.get() === 1)
+
+ // The parent is a ZippedWithIndexRDD whose start indices are already
computed,
+ // so they are reused directly and no extra job is submitted.
+ numJobs.set(0)
+ val rezipped = base.zipWithIndex()
+ sc.listenerBus.waitUntilEmpty()
+ assert(numJobs.get() === 0)
+ assert(rezipped.map(_._2).collect().toSeq === expected)
+
+ // map preserves partition sizes, so getAncestorWithSamePartitionSizes
walks up
+ // to `base` and reuses its start indices, again without a job.
+ sc.listenerBus.waitUntilEmpty()
+ numJobs.set(0)
+ val mapped = base.map { case (v, _) => v }.zipWithIndex()
+ sc.listenerBus.waitUntilEmpty()
+ assert(numJobs.get() === 0)
+ assert(mapped.map(_._2).collect().toSeq === expected)
+ } finally {
+ sc.removeSparkListener(listener)
+ }
+ }
+
+ test("zipWithIndex recomputes start indices when an intervening op changes
partition sizes") {
+ val numJobs = new AtomicInteger(0)
+ val listener = new SparkListener {
+ override def onJobStart(jobStart: SparkListenerJobStart): Unit = {
+ numJobs.incrementAndGet()
+ }
+ }
+ sc.addSparkListener(listener)
+ try {
+ val n = 10
+ val base = sc.parallelize(0 until n, 3).zipWithIndex()
+ sc.listenerBus.waitUntilEmpty()
+
+ // filter does not preserve partition sizes, so the ancestor walk stops
before
+ // `base`; a fresh counting job runs and the indices reflect the
filtered rows.
+ numJobs.set(0)
+ val filtered = base.filter { case (v, _) => v % 2 == 0 }.zipWithIndex()
+ sc.listenerBus.waitUntilEmpty()
+ assert(numJobs.get() === 1)
+ assert(filtered.map(_._2).collect().toSeq === (0 until 5).map(_.toLong))
+ } finally {
+ sc.removeSparkListener(listener)
+ }
+ }
+
test("zipWithUniqueId") {
val n = 10
val data = sc.parallelize(0 until n, 3)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
