perf: cache parsed plan data across a stage tasks

.github/workflows/pr_build_linux.yml

@@ -424,6 +424,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

.github/workflows/pr_build_macos.yml

@@ -197,6 +197,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

native/proto/src/proto/operator.proto

@@ -187,6 +187,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 {

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 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 = OperatorOuterClass.Operator.parseFrom(serializedPlan)
+      val basePlan = PlanDataInjector.parseBasePlan(serializedPlan)
       val injected =
         PlanDataInjector.injectPlanData(basePlan, commonByKey, partition.planDataByKey)
       PlanDataInjector.serializeOperator(injected)

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -366,11 +366,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,

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()
       }
     }
   }

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)

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()
   }
 }