From 843845940b4d8d3c4ab9efe9ec2a2851e2a33a04 Mon Sep 17 00:00:00 2001
From: Dustin Smith <dustin.william.smith@gmail.com>
Date: Tue, 1 Sep 2026 22:34:09 +0700
Subject: [PATCH 1/3] perf: cache parsed plan data across a stage tasks

Every task of a stage deserializes byte-identical plan bytes, yet each
one parsed the full operator tree, re-derived the scan source key by
stringifying its schema and filter lists, and re-parsed the scan
common message before injecting its partition data. Three bounded
per-executor caches now share that work: the parsed base plan keyed on
content, the parsed NativeScanCommon, and a source-key memo that hits
protobuf reference-identity fast path once the base plan is shared.
Injection itself stays per task, since partition data genuinely
differs, and the injected tree is never cached.

Per-task overhead drops from roughly 300us to 60us on a 100-column
scan plan and about 3x on a 1000-column plan. Cache misses compute
outside any lock so unrelated stages never serialize behind one parse,
and racing threads on a cold key adopt a single instance so reference
sharing holds.
---
 .../apache/spark/sql/comet/CometExecRDD.scala |   7 +-
 .../apache/spark/sql/comet/operators.scala    |  84 ++++++-
 .../sql/comet/PlanDataInjectorSuite.scala     | 213 ++++++++++++++++++
 3 files changed, 299 insertions(+), 5 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
index d9e0bf3a4c9..6953de743e3 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
@@ -29,7 +29,6 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.SerializableConfiguration
 
 import org.apache.comet.{CometExecIterator, CometRuntimeException, CometShuffleBlockIterator}
-import org.apache.comet.serde.OperatorOuterClass
 
 /**
  * Partition that carries per-partition planning data, avoiding closure capture of all partitions.
@@ -112,9 +111,11 @@ private[spark] class CometExecRDD(
         shuffleScanIndices,
         context)
 
-    // Only inject if we have per-partition planning data
+    // Only inject if we have per-partition planning data. The base plan bytes are identical
+    // for every partition of the stage, so the parsed tree is shared across this executor's
+    // tasks instead of being re-parsed per task.
     val actualPlan = if (commonByKey.nonEmpty) {
-      val basePlan = OperatorOuterClass.Operator.parseFrom(serializedPlan)
+      val basePlan = PlanDataInjector.parseBasePlan(serializedPlan)
       val injected =
         PlanDataInjector.injectPlanData(basePlan, commonByKey, partition.planDataByKey)
       PlanDataInjector.serializeOperator(injected)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index 3700e97642b..f6d1e06fa47 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -95,6 +95,56 @@ private[comet] trait PlanDataInjector {
  * Registry and utilities for injecting per-partition planning data into operator trees.
  */
 private[comet] object PlanDataInjector extends Logging {
+  import java.nio.ByteBuffer
+  import java.util.{LinkedHashMap, Map => JMap}
+
+  private[comet] final val maxCachedBasePlans = 16
+
+  // Every task of a stage deserializes its own byte-identical copy of the base plan, so
+  // without a cache an executor re-parses the same operator tree once per task. Parsed
+  // Operators are immutable, so one instance is safely shared across concurrent tasks.
+  // Keyed by content (ByteBuffer hashes/compares the bytes) since the arrays are distinct.
+  //
+  // Entries are whole parsed plan trees, so the entry count is what bounds executor memory:
+  // at most 16 recent stages' plans stay live, LRU-evicted as stages turn over. A stage
+  // rerun that misses after eviction simply re-parses.
+  private val basePlanCache = java.util.Collections.synchronizedMap(
+    new LinkedHashMap[ByteBuffer, Operator](4, 0.75f, true) {
+      override def removeEldestEntry(eldest: JMap.Entry[ByteBuffer, Operator]): Boolean = {
+        size() > maxCachedBasePlans
+      }
+    })
+
+  /**
+   * Look up `key`, computing and inserting the value on a miss. The computation runs outside any
+   * lock so unrelated misses never serialize behind each other; when two threads race the same
+   * cold key, the first insert wins and the loser adopts it, keeping the cached value
+   * reference-shared (which the sourceKey memo's identity fast path relies on).
+   */
+  private[comet] def cachedOrCompute[K, V](cache: JMap[K, V], key: K)(compute: => V): V = {
+    val cached = cache.get(key)
+    if (cached != null) {
+      cached
+    } else {
+      val computed = compute
+      cache.synchronized {
+        val winner = cache.get(key)
+        if (winner != null) {
+          winner
+        } else {
+          cache.put(key, computed)
+          computed
+        }
+      }
+    }
+  }
+
+  /**
+   * Parse a stage's base plan bytes, sharing the parsed tree across the executor's tasks. Falls
+   * back to a plain parse on eviction, so a stage rerun is always correct.
+   */
+  def parseBasePlan(bytes: Array[Byte]): Operator =
+    cachedOrCompute(basePlanCache, ByteBuffer.wrap(bytes))(Operator.parseFrom(bytes))
 
   // Registry of injectors for different operator types. The built-in injectors live in core.
   // Out-of-tree contribs (e.g. contrib-delta's `DeltaPlanDataInjector`) are discovered via the
@@ -328,6 +378,32 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector {
  * Injector for NativeScan operators.
  */
 private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
+  import java.nio.ByteBuffer
+  import java.util.{LinkedHashMap, Map => JMap}
+
+  private final val maxCacheEntries = 16
+
+  // Same rationale as IcebergPlanDataInjector's commonCache: the common bytes are identical
+  // for every partition of a stage, and parsing them dominates inject() for wide schemas.
+  private val commonCache = java.util.Collections.synchronizedMap(
+    new LinkedHashMap[ByteBuffer, OperatorOuterClass.NativeScanCommon](4, 0.75f, true) {
+      override def removeEldestEntry(
+          eldest: JMap.Entry[ByteBuffer, OperatorOuterClass.NativeScanCommon]): Boolean = {
+        size() > maxCacheEntries
+      }
+    })
+
+  // sourceKey stringifies the scan's schema and filter lists, which is the most expensive
+  // step of the whole per-task injection for wide schemas. Once the parsed base plan is
+  // shared across tasks the same common instance recurs, so protobuf's memoized hashCode
+  // and reference-equality fast path make this lookup O(1) after the first task.
+  private val keyCache = java.util.Collections.synchronizedMap(
+    new LinkedHashMap[OperatorOuterClass.NativeScanCommon, String](4, 0.75f, true) {
+      override def removeEldestEntry(
+          eldest: JMap.Entry[OperatorOuterClass.NativeScanCommon, String]): Boolean = {
+        size() > maxCacheEntries
+      }
+    })
 
   override val opStructCase: Operator.OpStructCase = Operator.OpStructCase.NATIVE_SCAN
 
@@ -336,7 +412,10 @@ private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
       op.getNativeScan.hasCommon &&
       !op.getNativeScan.hasFilePartition
 
-  override def getKey(op: Operator): Option[String] = Some(sourceKey(op.getNativeScan.getCommon))
+  override def getKey(op: Operator): Option[String] = {
+    val common = op.getNativeScan.getCommon
+    Some(PlanDataInjector.cachedOrCompute(keyCache, common)(sourceKey(common)))
+  }
 
   /**
    * The key under which a native scan's planning data is stored and looked up. Called on the
@@ -365,7 +444,8 @@ private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
       commonBytes: Array[Byte],
       partitionBytes: Array[Byte]): Operator = {
 
-    val common = OperatorOuterClass.NativeScanCommon.parseFrom(commonBytes)
+    val common = PlanDataInjector.cachedOrCompute(commonCache, ByteBuffer.wrap(commonBytes))(
+      OperatorOuterClass.NativeScanCommon.parseFrom(commonBytes))
     val partitionOnly = OperatorOuterClass.NativeScan.parseFrom(partitionBytes)
 
     // Build complete NativeScan with common fields + this partition's file list
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
index 91478b23491..40d8423cbde 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
@@ -143,6 +143,219 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     assert(IcebergPlanDataInjector.getKey(opA) == IcebergPlanDataInjector.getKey(opB))
   }
 
+  /** Builds an un-injected NativeScan operator: hasCommon, no file_partition. */
+  private def nativeScanOp(source: String, columnNames: Seq[String]): Operator = {
+    val common = nativeScanCommon(source, columnNames)
+    Operator
+      .newBuilder()
+      .setNativeScan(OperatorOuterClass.NativeScan.newBuilder().setCommon(common))
+      .build()
+  }
+
+  private def nativeScanCommon(
+      source: String,
+      columnNames: Seq[String]): OperatorOuterClass.NativeScanCommon = {
+    val builder = OperatorOuterClass.NativeScanCommon.newBuilder().setSource(source)
+    columnNames.foreach { name =>
+      builder.addRequiredSchema(
+        OperatorOuterClass.SparkStructField.newBuilder().setName(name).setNullable(true).build())
+    }
+    builder.build()
+  }
+
+  private def nativeScanPartitionBytes(filePath: String): Array[Byte] = {
+    OperatorOuterClass.NativeScan
+      .newBuilder()
+      .setFilePartition(
+        OperatorOuterClass.SparkFilePartition
+          .newBuilder()
+          .addPartitionedFile(
+            OperatorOuterClass.SparkPartitionedFile.newBuilder().setFilePath(filePath)))
+      .build()
+      .toByteArray
+  }
+
+  test("parseBasePlan shares one parsed Operator across byte-identical plans") {
+    val op = Operator
+      .newBuilder()
+      .setPlanId(10)
+      .addChildren(nativeScanOp("file:///cache-hit-tbl", Seq("a", "b")))
+      .build()
+    // Each Spark task deserializes its own copy of the task binary, so the bytes arrive as
+    // distinct arrays with identical content.
+    val bytes1 = op.toByteArray
+    val bytes2 = op.toByteArray
+    assert(!(bytes1 eq bytes2))
+
+    val parsed1 = PlanDataInjector.parseBasePlan(bytes1)
+    val parsed2 = PlanDataInjector.parseBasePlan(bytes2)
+
+    assert(parsed1 eq parsed2, "equal plan bytes should hit the cache, not re-parse")
+    assert(parsed1 == Operator.parseFrom(bytes1))
+  }
+
+  test("parseBasePlan keeps distinct plans separate") {
+    val opA = Operator
+      .newBuilder()
+      .setPlanId(20)
+      .addChildren(nativeScanOp("file:///distinct-tbl-a", Seq("a")))
+      .build()
+    val opB = Operator
+      .newBuilder()
+      .setPlanId(21)
+      .addChildren(nativeScanOp("file:///distinct-tbl-b", Seq("b")))
+      .build()
+
+    val parsedA = PlanDataInjector.parseBasePlan(opA.toByteArray)
+    val parsedB = PlanDataInjector.parseBasePlan(opB.toByteArray)
+
+    assert(parsedA == opA)
+    assert(parsedB == opB)
+    assert(parsedA != parsedB)
+  }
+
+  test("parseBasePlan re-parses correctly after eviction") {
+    val first = Operator
+      .newBuilder()
+      .setPlanId(30)
+      .addChildren(nativeScanOp("file:///evict-tbl-first", Seq("a")))
+      .build()
+    val firstBytes = first.toByteArray
+    val firstParsed = PlanDataInjector.parseBasePlan(firstBytes)
+
+    // Push enough distinct plans through to evict the first entry.
+    (0 until PlanDataInjector.maxCachedBasePlans).foreach { i =>
+      val filler = Operator
+        .newBuilder()
+        .setPlanId(1000 + i)
+        .addChildren(nativeScanOp(s"file:///evict-filler-$i", Seq("a")))
+        .build()
+      PlanDataInjector.parseBasePlan(filler.toByteArray)
+    }
+
+    val reParsed = PlanDataInjector.parseBasePlan(firstBytes)
+    assert(!(reParsed eq firstParsed), "the first plan should have been evicted")
+    assert(reParsed == first, "a rerun after eviction must still parse correctly")
+  }
+
+  test("parseBasePlan returns each thread the plan matching its bytes under concurrency") {
+    import java.util.concurrent.Executors
+    import scala.concurrent.{Await, ExecutionContext, Future}
+    import scala.concurrent.duration._
+
+    val plans = (0 until 4).map { i =>
+      Operator
+        .newBuilder()
+        .setPlanId(40 + i)
+        .addChildren(nativeScanOp(s"file:///concurrent-tbl-$i", Seq("a", "b")))
+        .build()
+    }
+    val pool = Executors.newFixedThreadPool(8)
+    implicit val ec: ExecutionContext = ExecutionContext.fromExecutor(pool)
+    try {
+      val checks = Future.sequence((0 until 64).map { i =>
+        val plan = plans(i % plans.size)
+        Future(PlanDataInjector.parseBasePlan(plan.toByteArray) == plan)
+      })
+      assert(Await.result(checks, 30.seconds).forall(identity))
+    } finally {
+      pool.shutdown()
+    }
+  }
+
+  test("parseBasePlan gives racing threads on a cold key the same instance") {
+    import java.util.concurrent.{CyclicBarrier, Executors}
+    import scala.concurrent.{Await, ExecutionContext, Future}
+    import scala.concurrent.duration._
+
+    val pool = Executors.newFixedThreadPool(2)
+    implicit val ec: ExecutionContext = ExecutionContext.fromExecutor(pool)
+    try {
+      (0 until 200).foreach { trial =>
+        val op = Operator
+          .newBuilder()
+          .setPlanId(5000 + trial)
+          .addChildren(nativeScanOp(s"file:///race-tbl-$trial", (0 until 64).map(i => s"c$i")))
+          .build()
+        val barrier = new CyclicBarrier(2)
+        val results = (0 until 2)
+          .map { _ =>
+            Future {
+              barrier.await()
+              PlanDataInjector.parseBasePlan(op.toByteArray)
+            }
+          }
+          .map(Await.result(_, 30.seconds))
+        // Whoever inserts first wins; the loser must adopt that instance, not its own parse,
+        // or downstream reference-identity sharing silently degrades.
+        assert(results(0) eq results(1), s"trial $trial: racing threads must share one instance")
+      }
+    } finally {
+      pool.shutdown()
+    }
+  }
+
+  test("NativeScan inject shares one parsed common across a stage's partitions") {
+    val scanOp = nativeScanOp("file:///shared-common-tbl", Seq("id", "v"))
+    val commonProto = nativeScanCommon("file:///shared-common-tbl", Seq("id", "v"))
+    val key = NativeScanPlanDataInjector.getKey(scanOp).get
+    // Distinct arrays with equal content, as two tasks of the same stage would hold them.
+    val commonBytes1 = commonProto.toByteArray
+    val commonBytes2 = commonProto.toByteArray
+
+    val injected1 = PlanDataInjector.injectPlanData(
+      scanOp,
+      Map(key -> commonBytes1),
+      Map(key -> nativeScanPartitionBytes("part-0.parquet")))
+    val injected2 = PlanDataInjector.injectPlanData(
+      scanOp,
+      Map(key -> commonBytes2),
+      Map(key -> nativeScanPartitionBytes("part-1.parquet")))
+
+    assert(
+      injected1.getNativeScan.getCommon eq injected2.getNativeScan.getCommon,
+      "equal common bytes should be parsed once and shared")
+    assert(injected1.getNativeScan.getCommon == commonProto)
+    // Each partition still gets its own file list.
+    val file1 = injected1.getNativeScan.getFilePartition.getPartitionedFile(0).getFilePath
+    val file2 = injected2.getNativeScan.getFilePartition.getPartitionedFile(0).getFilePath
+    assert(file1 == "part-0.parquet")
+    assert(file2 == "part-1.parquet")
+  }
+
+  test("NativeScan inject keeps different commons separate") {
+    val scanA = nativeScanOp("file:///separate-tbl-a", Seq("a"))
+    val scanB = nativeScanOp("file:///separate-tbl-b", Seq("b"))
+    val commonA = nativeScanCommon("file:///separate-tbl-a", Seq("a"))
+    val commonB = nativeScanCommon("file:///separate-tbl-b", Seq("b"))
+    val keyA = NativeScanPlanDataInjector.getKey(scanA).get
+    val keyB = NativeScanPlanDataInjector.getKey(scanB).get
+    assert(keyA != keyB)
+
+    val commonByKey = Map(keyA -> commonA.toByteArray, keyB -> commonB.toByteArray)
+    val partByKey = Map(
+      keyA -> nativeScanPartitionBytes("a.parquet"),
+      keyB -> nativeScanPartitionBytes("b.parquet"))
+
+    val injectedA = PlanDataInjector.injectPlanData(scanA, commonByKey, partByKey)
+    val injectedB = PlanDataInjector.injectPlanData(scanB, commonByKey, partByKey)
+
+    assert(injectedA.getNativeScan.getCommon == commonA)
+    assert(injectedB.getNativeScan.getCommon == commonB)
+  }
+
+  test("NativeScan getKey memo agrees with fresh derivation") {
+    val scanOp = nativeScanOp("file:///memo-tbl", Seq("id", "v", "w"))
+    // Same node twice (the shared-base-plan case), then an equal but distinct node.
+    val memoized = NativeScanPlanDataInjector.getKey(scanOp)
+    val again = NativeScanPlanDataInjector.getKey(scanOp)
+    val fresh =
+      NativeScanPlanDataInjector.getKey(nativeScanOp("file:///memo-tbl", Seq("id", "v", "w")))
+
+    assert(memoized == again)
+    assert(memoized == fresh)
+  }
+
   test(
     "self-join: scans sharing a metadataLocation but differing scan_hash_code inject their " +
       "own data, not each other's") {

From 36fa94d3ae84210da6ff0d3f9df5a31334f4b641 Mon Sep 17 00:00:00 2001
From: Dustin Smith <dustin.william.smith@gmail.com>
Date: Thu, 3 Sep 2026 11:21:58 +0700
Subject: [PATCH 2/3] perf: carry the scan key in the plan and scope prepared
 data to the base plan entry

The executor derived each native scan's lookup key by stringifying its
schema and filter lists, memoized in an LRU probed under the map
monitor, and parsed scan commons through a second per-scan LRU that a
single plan with 17 scans churned to zero reuse. The driver already
derives the key once in CometNativeScanExec, so carry it inside the
NativeScan proto and read it back on every injection path, including
the native shuffle writer's, which previously missed the fast paths
entirely.

Prepared commons now live inside the base plan's own cache entry
(scoped to the shuffleId on the shuffle-write path), so a plan and its
scans form one eviction unit. Entries pin the finalized common bytes
because scalar-subquery data filters resolve per execution: equal base
plan bytes do not guarantee equal finalized commons, and a stale entry
is replaced rather than served. The base plan cache keys on a stored
hash computed once per task outside the monitor instead of a raw
ByteBuffer that rescanned the bytes on every probe.
---
 native/proto/src/proto/operator.proto         |   5 +
 .../apache/spark/sql/comet/CometExecRDD.scala |   4 +-
 .../spark/sql/comet/CometNativeScanExec.scala |   8 +-
 .../shuffle/CometNativeShuffleWriter.scala    |   7 +-
 .../apache/spark/sql/comet/operators.scala    | 221 +++++++++++++-----
 .../comet/CometScanWithPlanDataSuite.scala    |   3 +-
 .../sql/comet/PlanDataInjectorSuite.scala     | 191 ++++++++++++---
 7 files changed, 346 insertions(+), 93 deletions(-)

diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
index 01bfa70c217..4c20a3993c1 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -184,6 +184,11 @@ message NativeScan {
 
   // Single partition's file list (injected at execution time)
   SparkFilePartition file_partition = 2;
+
+  // Key under which this scan's planning data is stored and looked up at execution time.
+  // Derived once on the driver (CometNativeScanExec) so executors read it back instead of
+  // re-deriving it per task. JVM-consumed only; the native side ignores it.
+  string source_key = 3;
 }
 
 message CsvScan {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
index 6953de743e3..10467f6c109 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
@@ -112,8 +112,8 @@ private[spark] class CometExecRDD(
         context)
 
     // Only inject if we have per-partition planning data. The base plan bytes are identical
-    // for every partition of the stage, so the parsed tree is shared across this executor's
-    // tasks instead of being re-parsed per task.
+    // for every partition of the stage, so the parsed tree and its prepared per-scan data are
+    // shared across this executor's tasks instead of being recomputed per task.
     val actualPlan = if (commonByKey.nonEmpty) {
       val basePlan = PlanDataInjector.parseBasePlan(serializedPlan)
       val injected =
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
index fbebc420046..d9f09f99cc2 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
@@ -365,11 +365,15 @@ object CometNativeScanExec {
       scan: CometScanExec): CometNativeScanExec = {
     // Generate unique key for this scan so PlanDataInjector can match common+partition data.
     // Multiple scans of same table with different projections/filters get different keys.
-    // Derived by the injector that will look it up, so the two sides cannot drift apart.
+    // Derived once here and embedded in the NativeScan proto, so executors (including the
+    // native shuffle writer) read it back instead of re-deriving it per task.
     val sourceKey = NativeScanPlanDataInjector.sourceKey(nativeOp.getNativeScan.getCommon)
+    val opWithKey = nativeOp.toBuilder
+      .setNativeScan(nativeOp.getNativeScan.toBuilder.setSourceKey(sourceKey))
+      .build()
 
     val batchScanExec = CometNativeScanExec(
-      nativeOp,
+      opWithKey,
       scanExec.relation,
       scanExec.output,
       scanExec.requiredSchema,
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index 89be100dbd3..7691cb16733 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -148,7 +148,12 @@ class CometNativeShuffleWriter[K, V](
       // in CometNativeShuffleInputRDD.getPartitions on the driver), not on the spec. The spec's
       // execContext.perPartitionByKey is emptied in prepareNativeShuffleDependency so the full
       // O(numPartitions) map stays out of the broadcast task binary.
-      PlanDataInjector.injectPlanData(
+      //
+      // The unified plan differs per task (output paths), so there is no base plan cache entry
+      // here; scan lookup rides the source keys the driver embedded in childNativeOp's scans,
+      // and prepared commons are shared across this shuffle's map tasks via the shuffleId.
+      PlanDataInjector.injectPlanDataForShuffle(
+        shuffleId,
         unifiedPlan,
         ctx.commonByKey,
         shuffleInputIter.planDataByKey)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index f6d1e06fa47..ff03106220a 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -82,35 +82,77 @@ private[comet] trait PlanDataInjector {
   def getKey(op: Operator): Option[String]
 
   /**
-   * Inject common + partition data into the operator node.
+   * Parse the partition-invariant common bytes into whatever form [[inject]] consumes.
+   * `injectPlanData` memoizes the result inside the base plan's cache entry, so a wide schema's
+   * common is prepared once per stage rather than once per task.
+   */
+  def prepareCommon(commonBytes: Array[Byte]): AnyRef
+
+  /**
+   * Inject common + partition data into the operator node. `preparedCommon` is what
+   * [[prepareCommon]] returned for this scan's common bytes.
    *
    * Implementations must return the node with its child list unchanged -- `injectPlanData` walks
    * the returned node's children, and relies on child reference identity to decide which
    * operators need rebuilding.
    */
-  def inject(op: Operator, commonBytes: Array[Byte], partitionBytes: Array[Byte]): Operator
+  def inject(op: Operator, preparedCommon: AnyRef, partitionBytes: Array[Byte]): Operator
 }
 
 /**
  * Registry and utilities for injecting per-partition planning data into operator trees.
  */
 private[comet] object PlanDataInjector extends Logging {
-  import java.nio.ByteBuffer
   import java.util.{LinkedHashMap, Map => JMap}
+  import java.util.concurrent.ConcurrentHashMap
 
   private[comet] final val maxCachedBasePlans = 16
 
+  /**
+   * Content key for the base plan cache. The hash is stored, computed once per task outside the
+   * cache monitor; a raw ByteBuffer key would rescan every byte on each probe under the lock.
+   */
+  private[comet] final class PlanKey(val bytes: Array[Byte]) {
+    private val hash: Int = java.util.Arrays.hashCode(bytes)
+    override def hashCode: Int = hash
+    override def equals(other: Any): Boolean = other match {
+      case that: PlanKey => (this eq that) || java.util.Arrays.equals(bytes, that.bytes)
+      case _ => false
+    }
+  }
+
+  /**
+   * A prepared common message together with the exact finalized bytes it was prepared from.
+   * Retaining the bytes costs one extra copy per scan entry; it is what makes the staleness check
+   * in prepareShared possible.
+   */
+  private[comet] final class PreparedCommon(val bytes: Array[Byte], val message: AnyRef)
+
+  /**
+   * A cached parsed base plan plus the per-scan data prepared for it. Prepared commons live
+   * inside their plan's cache entry, so the plan is the single eviction unit: a plan with any
+   * number of scans keeps them all while it stays cached, and can never churn another plan's
+   * scans out of a shared LRU.
+   */
+  private[comet] final class CachedPlanData(val plan: Operator) {
+    // Keyed by the scan's driver-computed source key. Entries pin the finalized common bytes
+    // they were prepared from: scalar-subquery data filters are appended after planning (see
+    // CometNativeScanExec.serializedPartitionData), so a byte-identical base plan can ship
+    // different finalized commons under the same key across executions.
+    private[comet] val preparedCommons = new ConcurrentHashMap[String, PreparedCommon]()
+  }
+
   // Every task of a stage deserializes its own byte-identical copy of the base plan, so
   // without a cache an executor re-parses the same operator tree once per task. Parsed
   // Operators are immutable, so one instance is safely shared across concurrent tasks.
-  // Keyed by content (ByteBuffer hashes/compares the bytes) since the arrays are distinct.
+  // Keyed by content (PlanKey hashes/compares the bytes) since the arrays are distinct.
   //
   // Entries are whole parsed plan trees, so the entry count is what bounds executor memory:
   // at most 16 recent stages' plans stay live, LRU-evicted as stages turn over. A stage
   // rerun that misses after eviction simply re-parses.
   private val basePlanCache = java.util.Collections.synchronizedMap(
-    new LinkedHashMap[ByteBuffer, Operator](4, 0.75f, true) {
-      override def removeEldestEntry(eldest: JMap.Entry[ByteBuffer, Operator]): Boolean = {
+    new LinkedHashMap[PlanKey, CachedPlanData](4, 0.75f, true) {
+      override def removeEldestEntry(eldest: JMap.Entry[PlanKey, CachedPlanData]): Boolean = {
         size() > maxCachedBasePlans
       }
     })
@@ -118,8 +160,8 @@ private[comet] object PlanDataInjector extends Logging {
   /**
    * Look up `key`, computing and inserting the value on a miss. The computation runs outside any
    * lock so unrelated misses never serialize behind each other; when two threads race the same
-   * cold key, the first insert wins and the loser adopts it, keeping the cached value
-   * reference-shared (which the sourceKey memo's identity fast path relies on).
+   * cold key, the first insert wins and the loser adopts it, so all tasks of a stage share one
+   * plan entry (and with it one set of prepared per-scan commons).
    */
   private[comet] def cachedOrCompute[K, V](cache: JMap[K, V], key: K)(compute: => V): V = {
     val cached = cache.get(key)
@@ -140,11 +182,13 @@ private[comet] object PlanDataInjector extends Logging {
   }
 
   /**
-   * Parse a stage's base plan bytes, sharing the parsed tree across the executor's tasks. Falls
-   * back to a plain parse on eviction, so a stage rerun is always correct.
+   * Parse a stage's base plan bytes, sharing the parsed tree and its prepared per-scan data
+   * across the executor's tasks. Falls back to a plain parse on eviction, so a stage rerun is
+   * always correct.
    */
-  def parseBasePlan(bytes: Array[Byte]): Operator =
-    cachedOrCompute(basePlanCache, ByteBuffer.wrap(bytes))(Operator.parseFrom(bytes))
+  def parseBasePlan(bytes: Array[Byte]): CachedPlanData =
+    cachedOrCompute(basePlanCache, new PlanKey(bytes))(
+      new CachedPlanData(Operator.parseFrom(bytes)))
 
   // Registry of injectors for different operator types. The built-in injectors live in core.
   // Out-of-tree contribs (e.g. contrib-delta's `DeltaPlanDataInjector`) are discovered via the
@@ -186,9 +230,52 @@ private[comet] object PlanDataInjector extends Logging {
    * reference rather than rebuilt; only the root-to-scan paths are rebuilt.
    */
   def injectPlanData(
+      op: Operator,
+      commonByKey: Map[String, Array[Byte]],
+      partitionByKey: Map[String, Array[Byte]]): Operator =
+    injectPlanData(op, commonByKey, partitionByKey, null)
+
+  /**
+   * Injects planning data into a cached base plan, memoizing each scan's prepared common inside
+   * the plan's own cache entry so it is prepared once per stage rather than once per task.
+   */
+  def injectPlanData(
+      cachedPlan: CachedPlanData,
+      commonByKey: Map[String, Array[Byte]],
+      partitionByKey: Map[String, Array[Byte]]): Operator =
+    injectPlanData(cachedPlan.plan, commonByKey, partitionByKey, cachedPlan.preparedCommons)
+
+  // The native shuffle writer's unified plan differs per task (output paths), so those plans
+  // never pass through parseBasePlan. Prepared commons for that path are scoped to the
+  // shuffleId instead: one shuffle stage's scans still share a single eviction unit.
+  private val shufflePreparedCommons = java.util.Collections.synchronizedMap(
+    new LinkedHashMap[Integer, ConcurrentHashMap[String, PreparedCommon]](4, 0.75f, true) {
+      override def removeEldestEntry(
+          eldest: JMap.Entry[Integer, ConcurrentHashMap[String, PreparedCommon]]): Boolean = {
+        size() > maxCachedBasePlans
+      }
+    })
+
+  /**
+   * Injects planning data into the native shuffle writer's per-task plan, sharing each scan's
+   * prepared common across the shuffle's map tasks. A shuffleId never spans executions, so the
+   * finalized common bytes under a key cannot change within one entry's lifetime.
+   */
+  def injectPlanDataForShuffle(
+      shuffleId: Int,
       op: Operator,
       commonByKey: Map[String, Array[Byte]],
       partitionByKey: Map[String, Array[Byte]]): Operator = {
+    val prepared = cachedOrCompute(shufflePreparedCommons, Integer.valueOf(shuffleId))(
+      new ConcurrentHashMap[String, PreparedCommon]())
+    injectPlanData(op, commonByKey, partitionByKey, prepared)
+  }
+
+  private def injectPlanData(
+      op: Operator,
+      commonByKey: Map[String, Array[Byte]],
+      partitionByKey: Map[String, Array[Byte]],
+      preparedCommons: ConcurrentHashMap[String, PreparedCommon]): Operator = {
 
     // O(1) by op kind, then a canInject confirm (which may inspect detail fields like `hasCommon`
     // / `!hasFilePartition`). Most operators in any tree are non-scan and skip the lookup body.
@@ -199,7 +286,8 @@ private[comet] object PlanDataInjector extends Logging {
             case Some(key) =>
               (commonByKey.get(key), partitionByKey.get(key)) match {
                 case (Some(commonBytes), Some(partitionBytes)) =>
-                  injector.inject(op, commonBytes, partitionBytes)
+                  val prepared = prepareShared(injector, key, commonBytes, preparedCommons)
+                  injector.inject(op, prepared, partitionBytes)
                 case _ =>
                   throw new CometRuntimeException(s"Missing planning data for key: $key")
               }
@@ -217,7 +305,7 @@ private[comet] object PlanDataInjector extends Logging {
     var i = 0
     while (i < numChildren) {
       val child = children.get(i)
-      val injectedChild = injectPlanData(child, commonByKey, partitionByKey)
+      val injectedChild = injectPlanData(child, commonByKey, partitionByKey, preparedCommons)
       if (injectedChild ne child) {
         if (builder == null) {
           builder = injectedOp.toBuilder
@@ -229,6 +317,32 @@ private[comet] object PlanDataInjector extends Logging {
     if (builder == null) injectedOp else builder.build()
   }
 
+  /**
+   * Prepared-common lookup scoped to the caller's cache entry, or a plain prepare when the caller
+   * has none. The byte comparison guards against a finalized common that changed under the same
+   * key -- scalar-subquery data filters resolve per execution -- so a stale entry is replaced,
+   * never served. Two overlapping executions alternating different finalized bytes under one key
+   * just alternate the slot: correct, only losing reuse for the overlap.
+   */
+  private def prepareShared(
+      injector: PlanDataInjector,
+      key: String,
+      commonBytes: Array[Byte],
+      preparedCommons: ConcurrentHashMap[String, PreparedCommon]): AnyRef = {
+    if (preparedCommons == null) {
+      injector.prepareCommon(commonBytes)
+    } else {
+      val hit = preparedCommons.get(key)
+      if (hit != null && java.util.Arrays.equals(hit.bytes, commonBytes)) {
+        hit.message
+      } else {
+        val prepared = new PreparedCommon(commonBytes, injector.prepareCommon(commonBytes))
+        preparedCommons.put(key, prepared)
+        prepared.message
+      }
+    }
+  }
+
   def serializeOperator(op: Operator): Array[Byte] = {
     val size = op.getSerializedSize
     val bytes = new Array[Byte](size)
@@ -348,21 +462,25 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector {
     Some(s"${common.getMetadataLocation}_${common.getScanHashCode}")
   }
 
-  override def inject(
-      op: Operator,
-      commonBytes: Array[Byte],
-      partitionBytes: Array[Byte]): Operator = {
-    val scan = op.getIcebergScan
-
-    // Cache the parsed common data to avoid deserializing on every partition
+  // Cache the parsed common data to avoid deserializing on every partition. Also serves the
+  // native shuffle writer's per-task plans, which have no base plan cache entry to memoize into.
+  override def prepareCommon(commonBytes: Array[Byte]): AnyRef = {
     val cacheKey = ByteBuffer.wrap(commonBytes)
-    val common = commonCache.synchronized {
+    commonCache.synchronized {
       Option(commonCache.get(cacheKey)).getOrElse {
         val parsed = OperatorOuterClass.IcebergScanCommon.parseFrom(commonBytes)
         commonCache.put(cacheKey, parsed)
         parsed
       }
     }
+  }
+
+  override def inject(
+      op: Operator,
+      preparedCommon: AnyRef,
+      partitionBytes: Array[Byte]): Operator = {
+    val scan = op.getIcebergScan
+    val common = preparedCommon.asInstanceOf[OperatorOuterClass.IcebergScanCommon]
 
     val tasksOnly = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes)
 
@@ -378,32 +496,6 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector {
  * Injector for NativeScan operators.
  */
 private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
-  import java.nio.ByteBuffer
-  import java.util.{LinkedHashMap, Map => JMap}
-
-  private final val maxCacheEntries = 16
-
-  // Same rationale as IcebergPlanDataInjector's commonCache: the common bytes are identical
-  // for every partition of a stage, and parsing them dominates inject() for wide schemas.
-  private val commonCache = java.util.Collections.synchronizedMap(
-    new LinkedHashMap[ByteBuffer, OperatorOuterClass.NativeScanCommon](4, 0.75f, true) {
-      override def removeEldestEntry(
-          eldest: JMap.Entry[ByteBuffer, OperatorOuterClass.NativeScanCommon]): Boolean = {
-        size() > maxCacheEntries
-      }
-    })
-
-  // sourceKey stringifies the scan's schema and filter lists, which is the most expensive
-  // step of the whole per-task injection for wide schemas. Once the parsed base plan is
-  // shared across tasks the same common instance recurs, so protobuf's memoized hashCode
-  // and reference-equality fast path make this lookup O(1) after the first task.
-  private val keyCache = java.util.Collections.synchronizedMap(
-    new LinkedHashMap[OperatorOuterClass.NativeScanCommon, String](4, 0.75f, true) {
-      override def removeEldestEntry(
-          eldest: JMap.Entry[OperatorOuterClass.NativeScanCommon, String]): Boolean = {
-        size() > maxCacheEntries
-      }
-    })
 
   override val opStructCase: Operator.OpStructCase = Operator.OpStructCase.NATIVE_SCAN
 
@@ -413,20 +505,23 @@ private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
       !op.getNativeScan.hasFilePartition
 
   override def getKey(op: Operator): Option[String] = {
-    val common = op.getNativeScan.getCommon
-    Some(PlanDataInjector.cachedOrCompute(keyCache, common)(sourceKey(common)))
+    val scan = op.getNativeScan
+    // The driver derives the key once and ships it inside the plan (CometNativeScanExec.apply),
+    // so no per-task derivation happens here; deriving from the common is only a fallback for
+    // plans built without one.
+    val transported = scan.getSourceKey
+    Some(if (transported.nonEmpty) transported else sourceKey(scan.getCommon))
   }
 
   /**
-   * The key under which a native scan's planning data is stored and looked up. Called on the
-   * driver by `CometNativeScanExec.apply` to store, and on the executor by [[getKey]] to look up
-   * \- both must derive the identical string from the same scan, so this is the single definition
-   * rather than two mirrored copies.
+   * The key under which a native scan's planning data is stored and looked up. Derived once on
+   * the driver by `CometNativeScanExec.apply`, which embeds it in the NativeScan proto so
+   * [[getKey]] reads the identical string back instead of re-deriving it.
    *
-   * Data filters are stripped of their `QueryContext` before hashing: the executor reads them
-   * back out of the interned plan (see `QueryContextInterner`) while the driver holds the
-   * un-interned form, so including the context encoding would make the two sides disagree. Only
-   * data filters can carry a context, so the other components are hashed as-is.
+   * Data filters are stripped of their `QueryContext` before hashing so the key is stable across
+   * interning (see `QueryContextInterner`): the executor-side fallback derivation sees the
+   * interned plan while the driver holds the un-interned form. Only data filters can carry a
+   * context, so the other components are hashed as-is.
    */
   private[comet] def sourceKey(common: OperatorOuterClass.NativeScanCommon): String = {
     val dataFilters = common.getDataFiltersList.asScala
@@ -439,13 +534,17 @@ private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
     s"${common.getSource}_${keyComponents.mkString("|").hashCode}"
   }
 
+  // Parsing wide-schema commons dominates inject(); injectPlanData memoizes the result in the
+  // base plan's cache entry so it runs once per stage on that path.
+  override def prepareCommon(commonBytes: Array[Byte]): AnyRef =
+    OperatorOuterClass.NativeScanCommon.parseFrom(commonBytes)
+
   override def inject(
       op: Operator,
-      commonBytes: Array[Byte],
+      preparedCommon: AnyRef,
       partitionBytes: Array[Byte]): Operator = {
 
-    val common = PlanDataInjector.cachedOrCompute(commonCache, ByteBuffer.wrap(commonBytes))(
-      OperatorOuterClass.NativeScanCommon.parseFrom(commonBytes))
+    val common = preparedCommon.asInstanceOf[OperatorOuterClass.NativeScanCommon]
     val partitionOnly = OperatorOuterClass.NativeScan.parseFrom(partitionBytes)
 
     // Build complete NativeScan with common fields + this partition's file list
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
index 0a3fc25a601..e5f7d08c05f 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
@@ -179,9 +179,10 @@ class TestStubPlanDataInjector extends PlanDataInjector {
   override def opStructCase: Operator.OpStructCase = Operator.OpStructCase.OPSTRUCT_NOT_SET
   override def canInject(op: Operator): Boolean = false
   override def getKey(op: Operator): Option[String] = None
+  override def prepareCommon(commonBytes: Array[Byte]): AnyRef = commonBytes
   override def inject(
       op: Operator,
-      commonBytes: Array[Byte],
+      preparedCommon: AnyRef,
       partitionBytes: Array[Byte]): Operator =
     op
 }
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
index 40d8423cbde..c74c4aa8b54 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
@@ -70,7 +70,10 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     val child = Operator.newBuilder().setPlanId(2).build()
     val root = Operator.newBuilder().setPlanId(1).addChildren(child).build()
 
-    val result = PlanDataInjector.injectPlanData(root, Map.empty, Map.empty)
+    val result = PlanDataInjector.injectPlanData(
+      root,
+      Map.empty[String, Array[Byte]],
+      Map.empty[String, Array[Byte]])
 
     assert(
       result eq root,
@@ -143,12 +146,19 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     assert(IcebergPlanDataInjector.getKey(opA) == IcebergPlanDataInjector.getKey(opB))
   }
 
-  /** Builds an un-injected NativeScan operator: hasCommon, no file_partition. */
+  /**
+   * Builds an un-injected NativeScan operator the way the driver ships it: hasCommon, no
+   * file_partition, source_key embedded (see CometNativeScanExec.apply).
+   */
   private def nativeScanOp(source: String, columnNames: Seq[String]): Operator = {
     val common = nativeScanCommon(source, columnNames)
     Operator
       .newBuilder()
-      .setNativeScan(OperatorOuterClass.NativeScan.newBuilder().setCommon(common))
+      .setNativeScan(
+        OperatorOuterClass.NativeScan
+          .newBuilder()
+          .setCommon(common)
+          .setSourceKey(NativeScanPlanDataInjector.sourceKey(common)))
       .build()
   }
 
@@ -191,7 +201,7 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     val parsed2 = PlanDataInjector.parseBasePlan(bytes2)
 
     assert(parsed1 eq parsed2, "equal plan bytes should hit the cache, not re-parse")
-    assert(parsed1 == Operator.parseFrom(bytes1))
+    assert(parsed1.plan == Operator.parseFrom(bytes1))
   }
 
   test("parseBasePlan keeps distinct plans separate") {
@@ -209,9 +219,9 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     val parsedA = PlanDataInjector.parseBasePlan(opA.toByteArray)
     val parsedB = PlanDataInjector.parseBasePlan(opB.toByteArray)
 
-    assert(parsedA == opA)
-    assert(parsedB == opB)
-    assert(parsedA != parsedB)
+    assert(parsedA.plan == opA)
+    assert(parsedB.plan == opB)
+    assert(parsedA.plan != parsedB.plan)
   }
 
   test("parseBasePlan re-parses correctly after eviction") {
@@ -235,7 +245,7 @@ class PlanDataInjectorSuite extends AnyFunSuite {
 
     val reParsed = PlanDataInjector.parseBasePlan(firstBytes)
     assert(!(reParsed eq firstParsed), "the first plan should have been evicted")
-    assert(reParsed == first, "a rerun after eviction must still parse correctly")
+    assert(reParsed.plan == first, "a rerun after eviction must still parse correctly")
   }
 
   test("parseBasePlan returns each thread the plan matching its bytes under concurrency") {
@@ -255,7 +265,7 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     try {
       val checks = Future.sequence((0 until 64).map { i =>
         val plan = plans(i % plans.size)
-        Future(PlanDataInjector.parseBasePlan(plan.toByteArray) == plan)
+        Future(PlanDataInjector.parseBasePlan(plan.toByteArray).plan == plan)
       })
       assert(Await.result(checks, 30.seconds).forall(identity))
     } finally {
@@ -299,22 +309,23 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     val scanOp = nativeScanOp("file:///shared-common-tbl", Seq("id", "v"))
     val commonProto = nativeScanCommon("file:///shared-common-tbl", Seq("id", "v"))
     val key = NativeScanPlanDataInjector.getKey(scanOp).get
-    // Distinct arrays with equal content, as two tasks of the same stage would hold them.
-    val commonBytes1 = commonProto.toByteArray
-    val commonBytes2 = commonProto.toByteArray
+    // Two tasks of the same stage: each deserializes its own copy of the plan and common bytes,
+    // and both resolve to the same cached base plan entry.
+    val task1 = PlanDataInjector.parseBasePlan(scanOp.toByteArray)
+    val task2 = PlanDataInjector.parseBasePlan(scanOp.toByteArray)
 
     val injected1 = PlanDataInjector.injectPlanData(
-      scanOp,
-      Map(key -> commonBytes1),
+      task1,
+      Map(key -> commonProto.toByteArray),
       Map(key -> nativeScanPartitionBytes("part-0.parquet")))
     val injected2 = PlanDataInjector.injectPlanData(
-      scanOp,
-      Map(key -> commonBytes2),
+      task2,
+      Map(key -> commonProto.toByteArray),
       Map(key -> nativeScanPartitionBytes("part-1.parquet")))
 
     assert(
       injected1.getNativeScan.getCommon eq injected2.getNativeScan.getCommon,
-      "equal common bytes should be parsed once and shared")
+      "equal common bytes should be prepared once per plan entry and shared")
     assert(injected1.getNativeScan.getCommon == commonProto)
     // Each partition still gets its own file list.
     val file1 = injected1.getNativeScan.getFilePartition.getPartitionedFile(0).getFilePath
@@ -323,6 +334,96 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     assert(file2 == "part-1.parquet")
   }
 
+  test("prepared scan data shares the base plan's eviction unit, not a per-scan budget") {
+    // A single plan with more scans than the base plan cache holds plans (17 > 16). All of the
+    // plan's prepared commons must be reused across tasks together; nothing may churn because
+    // the ownership unit is the plan entry, not a scan-count LRU.
+    val n = PlanDataInjector.maxCachedBasePlans + 1
+    val scans = (0 until n).map(i => nativeScanOp(s"file:///wide-plan-tbl-$i", Seq("a")))
+    val root = {
+      val builder = Operator.newBuilder().setPlanId(60)
+      scans.foreach(builder.addChildren)
+      builder.build()
+    }
+    val commonByKey = scans.map { s =>
+      NativeScanPlanDataInjector.getKey(s).get -> s.getNativeScan.getCommon.toByteArray
+    }.toMap
+    val partByKey = scans.zipWithIndex.map { case (s, i) =>
+      NativeScanPlanDataInjector.getKey(s).get -> nativeScanPartitionBytes(s"part-$i.parquet")
+    }.toMap
+    assert(commonByKey.size == n)
+
+    val first =
+      PlanDataInjector.injectPlanData(
+        PlanDataInjector.parseBasePlan(root.toByteArray),
+        commonByKey,
+        partByKey)
+    val second =
+      PlanDataInjector.injectPlanData(
+        PlanDataInjector.parseBasePlan(root.toByteArray),
+        commonByKey,
+        partByKey)
+
+    (0 until n).foreach { i =>
+      assert(
+        first.getChildren(i).getNativeScan.getCommon eq
+          second.getChildren(i).getNativeScan.getCommon,
+        s"scan $i must reuse the prepared common held by the plan's cache entry")
+    }
+  }
+
+  test("shuffle-path injection shares prepared commons across a shuffle's map tasks") {
+    // The native shuffle writer's unified plan differs per task, so it cannot share a base plan
+    // cache entry; prepared commons are scoped to the shuffleId instead.
+    val scanOp = nativeScanOp("file:///shuffle-share-tbl", Seq("id", "v"))
+    val key = NativeScanPlanDataInjector.getKey(scanOp).get
+    val common = scanOp.getNativeScan.getCommon
+
+    val task1 = PlanDataInjector.injectPlanDataForShuffle(
+      1234567,
+      scanOp,
+      Map(key -> common.toByteArray),
+      Map(key -> nativeScanPartitionBytes("map-0.parquet")))
+    val task2 = PlanDataInjector.injectPlanDataForShuffle(
+      1234567,
+      scanOp,
+      Map(key -> common.toByteArray),
+      Map(key -> nativeScanPartitionBytes("map-1.parquet")))
+
+    assert(
+      task1.getNativeScan.getCommon eq task2.getNativeScan.getCommon,
+      "one shuffle's map tasks must share the prepared common, not re-parse it")
+    assert(task1.getNativeScan.getCommon == common)
+  }
+
+  test("a changed finalized common under the same key is replaced, not served stale") {
+    // Scalar-subquery data filters are appended to the finalized common after planning
+    // (CometNativeScanExec.serializedPartitionData), so a byte-identical base plan can ship
+    // different finalized commons under the same transported key across executions.
+    val scanOp = nativeScanOp("file:///stale-common-tbl", Seq("id"))
+    val key = NativeScanPlanDataInjector.getKey(scanOp).get
+    val baseCommon = scanOp.getNativeScan.getCommon
+    val finalizedCommon = baseCommon.toBuilder
+      .addDataFilters(org.apache.comet.serde.ExprOuterClass.Expr.newBuilder())
+      .build()
+    assert(baseCommon != finalizedCommon)
+
+    val cached = PlanDataInjector.parseBasePlan(scanOp.toByteArray)
+    val firstRun = PlanDataInjector.injectPlanData(
+      cached,
+      Map(key -> baseCommon.toByteArray),
+      Map(key -> nativeScanPartitionBytes("run-1.parquet")))
+    val secondRun = PlanDataInjector.injectPlanData(
+      cached,
+      Map(key -> finalizedCommon.toByteArray),
+      Map(key -> nativeScanPartitionBytes("run-2.parquet")))
+
+    assert(firstRun.getNativeScan.getCommon == baseCommon)
+    assert(
+      secondRun.getNativeScan.getCommon == finalizedCommon,
+      "changed common bytes under the same key must be re-prepared, never served stale")
+  }
+
   test("NativeScan inject keeps different commons separate") {
     val scanA = nativeScanOp("file:///separate-tbl-a", Seq("a"))
     val scanB = nativeScanOp("file:///separate-tbl-b", Seq("b"))
@@ -344,16 +445,54 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     assert(injectedB.getNativeScan.getCommon == commonB)
   }
 
-  test("NativeScan getKey memo agrees with fresh derivation") {
-    val scanOp = nativeScanOp("file:///memo-tbl", Seq("id", "v", "w"))
-    // Same node twice (the shared-base-plan case), then an equal but distinct node.
-    val memoized = NativeScanPlanDataInjector.getKey(scanOp)
-    val again = NativeScanPlanDataInjector.getKey(scanOp)
-    val fresh =
-      NativeScanPlanDataInjector.getKey(nativeScanOp("file:///memo-tbl", Seq("id", "v", "w")))
+  test("NativeScan getKey reads the driver-computed source key from the plan") {
+    val common = nativeScanCommon("file:///transported-tbl", Seq("id", "v"))
+    val op = Operator
+      .newBuilder()
+      .setNativeScan(
+        OperatorOuterClass.NativeScan
+          .newBuilder()
+          .setCommon(common)
+          .setSourceKey("driver-key"))
+      .build()
+
+    assert(NativeScanPlanDataInjector.getKey(op).contains("driver-key"))
+  }
+
+  test("NativeScan getKey derives the key only when the plan carries none") {
+    val common = nativeScanCommon("file:///fallback-tbl", Seq("id", "v", "w"))
+    val op = Operator
+      .newBuilder()
+      .setNativeScan(OperatorOuterClass.NativeScan.newBuilder().setCommon(common))
+      .build()
+
+    val derived = NativeScanPlanDataInjector.sourceKey(common)
+    assert(NativeScanPlanDataInjector.getKey(op).contains(derived))
+  }
+
+  test("direct injection without a cached base plan looks up by the transported key") {
+    // CometNativeShuffleWriter injects into a per-task plan built around spec.childNativeOp
+    // without going through parseBasePlan, so the lookup must ride the transported key alone.
+    val common = nativeScanCommon("file:///shuffle-tbl", Seq("id"))
+    val op = Operator
+      .newBuilder()
+      .setNativeScan(
+        OperatorOuterClass.NativeScan
+          .newBuilder()
+          .setCommon(common)
+          .setSourceKey("shuffle-key"))
+      .build()
 
-    assert(memoized == again)
-    assert(memoized == fresh)
+    val injected = PlanDataInjector.injectPlanData(
+      op,
+      Map("shuffle-key" -> common.toByteArray),
+      Map("shuffle-key" -> nativeScanPartitionBytes("shuffled.parquet")))
+
+    assert(injected.getNativeScan.getCommon == common)
+    assert(
+      injected.getNativeScan.getFilePartition
+        .getPartitionedFile(0)
+        .getFilePath == "shuffled.parquet")
   }
 
   test(

From f24e4fcc6bc61d06b82b66450898e6197db82c36 Mon Sep 17 00:00:00 2001
From: Dustin Smith <dustin.william.smith@gmail.com>
Date: Thu, 3 Sep 2026 19:20:28 +0700
Subject: [PATCH 3/3] fix: release prepared shuffle data when a shuffle is
 unregistered or the manager stops

The shuffle-scoped prepared-commons store is a JVM singleton keyed by
shuffleId, and shuffle ids restart at zero for every SparkContext, so a
local or embedded caller that stops and recreates its context kept
stacking new scan keys under ids the previous context had used. Both
CometShuffleManager and CometCelebornShuffleManager now drop a shuffle's
entry in unregisterShuffle (reached from Spark's ContextCleaner via
BlockManagerStorageEndpoint) and clear the store in stop(), so a
recreated context starts empty and long-lived contexts release each
shuffle's data with the shuffle.
---
 .github/workflows/pr_build_linux.yml          |   1 +
 .github/workflows/pr_build_macos.yml          |   1 +
 .../shuffle/CometCelebornShuffleManager.scala |   3 +
 .../shuffle/CometShuffleManager.scala         |   5 +-
 .../apache/spark/sql/comet/operators.scala    |  17 +++
 ...lanDataInjectorShuffleLifecycleSuite.scala | 124 ++++++++++++++++++
 .../sql/comet/PlanDataInjectorSuite.scala     |  49 +++++++
 .../CometCelebornShuffleManagerSuite.scala    |  51 +++++++
 8 files changed, 250 insertions(+), 1 deletion(-)
 create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorShuffleLifecycleSuite.scala

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 2426db65fe4..c8fef9ea766 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -422,6 +422,7 @@ jobs:
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite
+              org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite
               org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite
               org.apache.spark.sql.comet.CometDecimalPromotionSuite
               org.apache.spark.sql.comet.CometScanWithPlanDataSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 6e4b50a8617..40e0bbd950a 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -195,6 +195,7 @@ jobs:
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite
+              org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite
               org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite
               org.apache.spark.sql.comet.CometDecimalPromotionSuite
               org.apache.spark.sql.comet.CometScanWithPlanDataSuite
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
index 88a9413e1dc..dccff7cbc6a 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
@@ -31,6 +31,7 @@ import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext}
 import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}
 import org.apache.spark.scheduler.OutputCommitCoordinator
 import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.sql.comet.PlanDataInjector
 import org.apache.spark.util.RpcUtils
 
 import org.apache.comet.CometConf
@@ -246,6 +247,7 @@ class CometCelebornShuffleManager private[shuffle] (
     if (isDriver) {
       Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId))
     }
+    PlanDataInjector.releasePreparedShuffle(shuffleId)
     backend.unregisterShuffle(shuffleId)
   }
 
@@ -262,6 +264,7 @@ class CometCelebornShuffleManager private[shuffle] (
         ownedNativeClients.keySet().asScala.foreach(CelebornShufflePusherFactory.releaseClient)
         ownedNativeClients.clear()
         nativeShuffleClients.clear()
+        PlanDataInjector.releaseAllPreparedShuffles()
       }
     }
   }
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
index 398ed66b6a5..d5976a46369 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
@@ -31,6 +31,7 @@ import org.apache.spark.internal.{config, Logging}
 import org.apache.spark.shuffle._
 import org.apache.spark.shuffle.api.ShuffleExecutorComponents
 import org.apache.spark.shuffle.sort.{BypassMergeSortShuffleHandle, SerializedShuffleHandle, SortShuffleManager, SortShuffleWriter}
+import org.apache.spark.sql.comet.PlanDataInjector
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.util.collection.OpenHashSet
 
@@ -282,12 +283,14 @@ class CometShuffleManager(conf: SparkConf) extends ShuffleManager with Logging {
         shuffleBlockResolver.removeDataByMap(shuffleId, mapTaskId)
       }
     }
+    PlanDataInjector.releasePreparedShuffle(shuffleId)
     true
   }
 
   /** Shut down this ShuffleManager. */
   override def stop(): Unit = {
-    shuffleBlockResolver.stop()
+    try shuffleBlockResolver.stop()
+    finally PlanDataInjector.releaseAllPreparedShuffles()
   }
 }
 
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index ff03106220a..2548318c9ee 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -271,6 +271,23 @@ private[comet] object PlanDataInjector extends Logging {
     injectPlanData(op, commonByKey, partitionByKey, prepared)
   }
 
+  // A shuffle's prepared data dies with the shuffle, and shuffle ids restart at zero for every
+  // SparkContext in the JVM, so a recreated context would otherwise keep stacking new scan keys
+  // under ids the last context already used. The shuffle managers call these from
+  // unregisterShuffle and stop.
+  private[comet] def releasePreparedShuffle(shuffleId: Int): Unit =
+    shufflePreparedCommons.remove(Integer.valueOf(shuffleId))
+
+  private[comet] def releaseAllPreparedShuffles(): Unit = shufflePreparedCommons.clear()
+
+  /** Test-only view: each cached shuffle id with the scan keys prepared under it. */
+  private[comet] def preparedShuffleSnapshot: Map[Int, Set[String]] =
+    shufflePreparedCommons.synchronized {
+      shufflePreparedCommons.asScala.map { case (id, prepared) =>
+        id.intValue() -> prepared.keySet().asScala.toSet
+      }.toMap
+    }
+
   private def injectPlanData(
       op: Operator,
       commonByKey: Map[String, Array[Byte]],
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorShuffleLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorShuffleLifecycleSuite.scala
new file mode 100644
index 00000000000..c7fb8f16a56
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorShuffleLifecycleSuite.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import scala.concurrent.duration.DurationInt
+
+import org.apache.spark.sql.{CometTestBase, SparkSession}
+import org.apache.spark.sql.functions.col
+import org.apache.spark.util.Utils
+
+import org.apache.comet.CometConf
+
+/**
+ * End-to-end lifecycle of the shuffle-scoped prepared scan data held by [[PlanDataInjector]]:
+ * Spark's shuffle cleanup releases one shuffle's entry, and stopping the SparkContext releases
+ * them all, so a context recreated in the same JVM starts from an empty store even though its
+ * shuffle ids restart at zero.
+ *
+ * The recreated-context test stops the suite's SparkContext, so it runs last and nothing else may
+ * follow it.
+ */
+class PlanDataInjectorShuffleLifecycleSuite extends CometTestBase {
+
+  private def withNativeShuffle[T](session: SparkSession)(f: => T): T = {
+    val keys = Seq(
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_SHUFFLE_MODE.key -> "native")
+    val previous = keys.map { case (k, _) => k -> session.conf.getOption(k) }
+    keys.foreach { case (k, v) => session.conf.set(k, v) }
+    try f
+    finally
+      previous.foreach {
+        case (k, Some(v)) => session.conf.set(k, v)
+        case (k, None) => session.conf.unset(k)
+      }
+  }
+
+  /** Runs a native shuffle fed directly by a native Parquet scan and returns its scan keys. */
+  private def runScanFusedShuffle(session: SparkSession, path: String): Set[String] = {
+    session
+      .range(0, 1000, 1, numPartitions = 4)
+      .selectExpr("id AS _1", "CAST(id AS STRING) AS _2")
+      .write
+      .parquet(path)
+    withNativeShuffle(session) {
+      val before = PlanDataInjector.preparedShuffleSnapshot
+      val df = session.read.parquet(path).repartition(5, col("_1"))
+      assert(df.count() == 1000)
+      val added = PlanDataInjector.preparedShuffleSnapshot.filterNot { case (id, keys) =>
+        before.get(id).contains(keys)
+      }
+      assert(added.size == 1, s"expected one new shuffle entry, saw $added")
+      val keys = added.values.head
+      assert(keys.nonEmpty, "the native scan must have been prepared under the shuffle's id")
+      keys
+    }
+  }
+
+  test("Spark's shuffle cleanup releases the shuffle's prepared scan data") {
+    PlanDataInjector.releaseAllPreparedShuffles()
+    withTempDir { dir =>
+      val keys = runScanFusedShuffle(spark, new java.io.File(dir, "cleanup.parquet").toString)
+      // The DataFrame is out of scope here; once the ShuffleDependency is collected, the
+      // ContextCleaner asks every block manager to remove the shuffle, which reaches
+      // CometShuffleManager.unregisterShuffle.
+      eventually(timeout(30.seconds), interval(1.second)) {
+        System.gc()
+        val live = PlanDataInjector.preparedShuffleSnapshot.values.flatten.toSet
+        assert(
+          (live & keys).isEmpty,
+          s"the collected shuffle's scan data should be gone, still holds ${live & keys}")
+      }
+    }
+  }
+
+  test("a recreated SparkContext starts from an empty shuffle store") {
+    PlanDataInjector.releaseAllPreparedShuffles()
+    // Not withTempDir: its task-drain check needs the suite's context, which this test stops.
+    val dir = Utils.createTempDir()
+    try {
+      val firstKeys =
+        runScanFusedShuffle(spark, new java.io.File(dir, "first-context.parquet").toString)
+      assert(PlanDataInjector.preparedShuffleSnapshot.values.flatten.toSet == firstKeys)
+
+      spark.stop()
+      assert(
+        PlanDataInjector.preparedShuffleSnapshot.isEmpty,
+        "stopping the context must release every shuffle's prepared data")
+
+      val second = createSparkSession
+      try {
+        val secondKeys =
+          runScanFusedShuffle(second, new java.io.File(dir, "second-context.parquet").toString)
+        val snapshot = PlanDataInjector.preparedShuffleSnapshot
+        assert(snapshot.size == 1, s"the new context should own the only entry, saw $snapshot")
+        assert(
+          snapshot.values.head == secondKeys,
+          "a shuffle id reused by the new context must not carry the old context's scans")
+      } finally {
+        second.stop()
+      }
+    } finally {
+      Utils.deleteRecursively(dir)
+    }
+  }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
index c74c4aa8b54..b32b721e7f2 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
@@ -21,6 +21,9 @@ package org.apache.spark.sql.comet
 
 import org.scalatest.funsuite.AnyFunSuite
 
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager
+
 import org.apache.comet.serde.OperatorOuterClass
 import org.apache.comet.serde.OperatorOuterClass.Operator
 
@@ -396,6 +399,52 @@ class PlanDataInjectorSuite extends AnyFunSuite {
     assert(task1.getNativeScan.getCommon == common)
   }
 
+  /** Injects one scan under `shuffleId` the way a map task would, returning its scan key. */
+  private def injectShuffleScan(shuffleId: Int, source: String): String = {
+    val scanOp = nativeScanOp(source, Seq("id"))
+    val key = NativeScanPlanDataInjector.getKey(scanOp).get
+    PlanDataInjector.injectPlanDataForShuffle(
+      shuffleId,
+      scanOp,
+      Map(key -> scanOp.getNativeScan.getCommon.toByteArray),
+      Map(key -> nativeScanPartitionBytes("map-0.parquet")))
+    key
+  }
+
+  test("recreated context does not accumulate prepared commons under reused shuffle ids") {
+    // Shuffle ids restart at zero for every SparkContext in a JVM, so a local or embedded caller
+    // that stops and recreates its context reuses ids the previous context already cached under.
+    // The manager's stop() is the boundary where the old context's prepared data must go.
+    PlanDataInjector.releaseAllPreparedShuffles()
+    val firstContext = new CometShuffleManager(new SparkConf(false))
+    val a = injectShuffleScan(0, "file:///recreated-ctx-a")
+    val b = injectShuffleScan(0, "file:///recreated-ctx-b")
+    assert(PlanDataInjector.preparedShuffleSnapshot(0) == Set(a, b))
+
+    firstContext.stop()
+
+    new CometShuffleManager(new SparkConf(false))
+    val c = injectShuffleScan(0, "file:///recreated-ctx-c")
+    val d = injectShuffleScan(0, "file:///recreated-ctx-d")
+    assert(
+      PlanDataInjector.preparedShuffleSnapshot(0) == Set(c, d),
+      "shuffle 0 must hold only the new context's scans, not the stopped context's as well")
+  }
+
+  test("unregisterShuffle releases only that shuffle's prepared commons") {
+    PlanDataInjector.releaseAllPreparedShuffles()
+    val manager = new CometShuffleManager(new SparkConf(false))
+    val gone = injectShuffleScan(7, "file:///unregister-gone")
+    val kept = injectShuffleScan(8, "file:///unregister-kept")
+    assert(PlanDataInjector.preparedShuffleSnapshot == Map(7 -> Set(gone), 8 -> Set(kept)))
+
+    manager.unregisterShuffle(7)
+
+    assert(
+      PlanDataInjector.preparedShuffleSnapshot == Map(8 -> Set(kept)),
+      "unregistering shuffle 7 must drop exactly its entry")
+  }
+
   test("a changed finalized common under the same key is replaced, not served stale") {
     // Scalar-subquery data filters are appended to the finalized common after planning
     // (CometNativeScanExec.serializedPartitionData), so a byte-identical base plan can ship
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
index f12a8b9f00e..73edb8880a4 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
@@ -26,6 +26,9 @@ import org.scalatest.funsuite.AnyFunSuite
 import org.apache.spark.{ExecutorLostFailure, ShuffleDependency, SparkConf, TaskContext, TaskEndReason, UnknownReason}
 import org.apache.spark.scheduler.OutputCommitCoordinator
 import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.sql.comet.{NativeScanPlanDataInjector, PlanDataInjector}
+
+import org.apache.comet.serde.OperatorOuterClass
 
 class CometCelebornShuffleManagerSuite extends AnyFunSuite {
 
@@ -573,6 +576,54 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
     assert(backend.stopped)
   }
 
+  /** Caches one prepared scan under `shuffleId` as a native map task would. */
+  private def prepareShuffleScan(shuffleId: Int, source: String): String = {
+    val common = OperatorOuterClass.NativeScanCommon.newBuilder().setSource(source).build()
+    val key = NativeScanPlanDataInjector.sourceKey(common)
+    val scanOp = OperatorOuterClass.Operator
+      .newBuilder()
+      .setNativeScan(
+        OperatorOuterClass.NativeScan.newBuilder().setCommon(common).setSourceKey(key))
+      .build()
+    val partition = OperatorOuterClass.NativeScan
+      .newBuilder()
+      .setFilePartition(OperatorOuterClass.SparkFilePartition.newBuilder())
+      .build()
+      .toByteArray
+    PlanDataInjector.injectPlanDataForShuffle(
+      shuffleId,
+      scanOp,
+      Map(key -> common.toByteArray),
+      Map(key -> partition))
+    key
+  }
+
+  test("unregisterShuffle releases that shuffle's prepared scan data") {
+    PlanDataInjector.releaseAllPreparedShuffles()
+    val composite = manager(new RecordingShuffleManager)
+    val gone = prepareShuffleScan(3, "s3://celeborn/unregister-gone")
+    val kept = prepareShuffleScan(4, "s3://celeborn/unregister-kept")
+    assert(PlanDataInjector.preparedShuffleSnapshot == Map(3 -> Set(gone), 4 -> Set(kept)))
+
+    composite.unregisterShuffle(3)
+
+    assert(PlanDataInjector.preparedShuffleSnapshot == Map(4 -> Set(kept)))
+  }
+
+  test("stop releases every shuffle's prepared scan data") {
+    // A recreated SparkContext restarts shuffle ids at zero, so whatever the stopped context
+    // cached under those ids must not survive the manager that owned it.
+    PlanDataInjector.releaseAllPreparedShuffles()
+    val composite = manager(new RecordingShuffleManager)
+    prepareShuffleScan(0, "s3://celeborn/stop-a")
+    prepareShuffleScan(1, "s3://celeborn/stop-b")
+    assert(PlanDataInjector.preparedShuffleSnapshot.keySet == Set(0, 1))
+
+    composite.stop()
+
+    assert(PlanDataInjector.preparedShuffleSnapshot.isEmpty)
+  }
+
   test("registration failures retain their original exception without local fallback") {
     val backend = new RecordingShuffleManager
     val expected = new IllegalStateException("remote shuffle registration failed")