perf: avoid rebuilding untouched operators in PlanDataInjector.injectPlanData

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

@@ -83,7 +83,13 @@ private[comet] trait PlanDataInjector {
   /** Extract the key used to look up planning data for this operator. */
   def getKey(op: Operator): Option[String]
 
-  /** Inject common + partition data into the operator node. */
+  /**
+   * Inject common + partition data into the operator node.
+   *
+   * 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
 }
 
@@ -132,40 +138,51 @@ private[comet] object PlanDataInjector extends Logging {
    *
    * Supports joins over multiple tables by matching each operator with its corresponding data
    * based on a key (e.g., metadata_location for Iceberg).
+   *
+   * Operators are immutable protobuf messages, so any subtree needing no injection is returned by
+   * 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 = {
-    val builder = op.toBuilder
 
     // 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.
-    injectorsByKind.get(op.getOpStructCase) match {
+    val injectedOp = injectorsByKind.get(op.getOpStructCase) match {
       case Some(injector) if injector.canInject(op) =>
         injector.getKey(op) match {
           case Some(key) =>
             (commonByKey.get(key), partitionByKey.get(key)) match {
               case (Some(commonBytes), Some(partitionBytes)) =>
-                val injectedOp = injector.inject(op, commonBytes, partitionBytes)
-                // Copy the injected operator's fields to our builder
-                builder.clear()
-                builder.mergeFrom(injectedOp)
+                injector.inject(op, commonBytes, partitionBytes)
               case _ =>
                 throw new CometRuntimeException(s"Missing planning data for key: $key")
             }
-          case None =>
+          case None => op
         }
-      case _ =>
-    }
-
-    // Recursively process children
-    builder.clearChildren()
-    op.getChildrenList.asScala.foreach { child =>
-      builder.addChildren(injectPlanData(child, commonByKey, partitionByKey))
+      case _ => op
+    }
+
+    // Recursively process children, rebuilding this node only if one of them actually changed.
+    // Injectors preserve children, so `injectedOp` has the same child list as `op` either way.
+    // The builder is created on the first changed child, so unchanged nodes allocate nothing.
+    val children = injectedOp.getChildrenList
+    val numChildren = children.size()
+    var builder: Operator.Builder = null
+    var i = 0
+    while (i < numChildren) {
+      val child = children.get(i)
+      val injectedChild = injectPlanData(child, commonByKey, partitionByKey)
+      if (injectedChild ne child) {
+        if (builder == null) {
+          builder = injectedOp.toBuilder
+        }
+        builder.setChildren(i, injectedChild)
+      }
+      i += 1
     }
-
-    builder.build()
+    if (builder == null) injectedOp else builder.build()
   }
 
   def serializeOperator(op: Operator): Array[Byte] = {

spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala

@@ -72,7 +72,47 @@ class PlanDataInjectorSuite extends AnyFunSuite {
 
     val result = PlanDataInjector.injectPlanData(root, Map.empty, Map.empty)
 
-    assert(result == root, "non-scan operator tree should be returned unchanged")
+    assert(
+      result eq root,
+      "a tree with nothing to inject should be returned by reference, not rebuilt")
+  }
+
+  test("injectPlanData rebuilds only the path to the injected scan") {
+    // Operators are immutable protobuf messages, so subtrees that need no injection are shared.
+    val scanOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111)
+    val (commonBytes, partitionBytes) =
+      icebergPlanData(
+        "s3://table/metadata/v1.json",
+        scanHashCode = 111,
+        columnNames = Seq("id", "v"),
+        dataFilePath = "data.parquet")
+    val key = IcebergPlanDataInjector.getKey(scanOp).get
+
+    val filter = Operator.newBuilder().setPlanId(2).addChildren(scanOp).build()
+    val untouchedSibling = Operator
+      .newBuilder()
+      .setPlanId(3)
+      .addChildren(Operator.newBuilder().setPlanId(4).build())
+      .build()
+    val root = Operator
+      .newBuilder()
+      .setPlanId(1)
+      .addChildren(filter)
+      .addChildren(untouchedSibling)
+      .build()
+
+    val result =
+      PlanDataInjector.injectPlanData(root, Map(key -> commonBytes), Map(key -> partitionBytes))
+
+    assert(
+      result.getChildren(1) eq untouchedSibling,
+      "a sibling subtree with no injectable scan should be shared, not rebuilt")
+    val injectedScan = result.getChildren(0).getChildren(0)
+    assert(injectedScan.getIcebergScan.getCommon.getRequiredSchemaCount == 2)
+    assert(injectedScan.getIcebergScan.getFileScanTasks(0).getDataFilePath == "data.parquet")
+    // Everything outside the injected scan is preserved verbatim.
+    assert(result.getPlanId == 1)
+    assert(result.getChildren(0).getPlanId == 2)
   }
 
   test("each registered injector is reachable by its opStructCase") {