fixup

apache · Oct 9, 2024 · 6edc04c · 6edc04c
1 parent 60c46e8
commit 6edc04c
Showing 1 changed file with 3 additions and 15 deletions.
diff --git a/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala b/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala
@@ -526,27 +526,15 @@ class Spark35Shims extends SparkShims {
     Seq(expr.srcArrayExpr, expr.posExpr, expr.itemExpr, Literal(expr.legacyNegativeIndex))
   }
 
-  override def withOperatorIdMap[T](idMap: java.util.Map[QueryPlan[_], Int])(body: => T): T = {
-    val prevIdMap = QueryPlan.localIdMap.get()
-    try {
-      QueryPlan.localIdMap.set(idMap)
-      body
-    } finally {
-      QueryPlan.localIdMap.set(prevIdMap)
-    }
-  }
-
   override def getOperatorId(plan: QueryPlan[_]): Option[Int] = {
-    Option(QueryPlan.localIdMap.get().get(plan))
+    plan.getTagValue(QueryPlan.OP_ID_TAG)
   }
 
   override def setOperatorId(plan: QueryPlan[_], opId: Int): Unit = {
-    val map = QueryPlan.localIdMap.get()
-    assert(!map.containsKey(plan))
-    map.put(plan, opId)
+    plan.setTagValue(QueryPlan.OP_ID_TAG, opId)
   }
 
   override def unsetOperatorId(plan: QueryPlan[_]): Unit = {
-    QueryPlan.localIdMap.get().remove(plan)
+    plan.unsetTagValue(QueryPlan.OP_ID_TAG)
   }
 }