[SPARK-33427][SQL][FOLLOWUP] Put key and value into IdentityHashMap sequantially

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SubExprEvaluationRuntime.scala

@@ -18,8 +18,6 @@ package org.apache.spark.sql.catalyst.expressions
 
 import java.util.IdentityHashMap
 
-import scala.collection.JavaConverters._
-
 import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache}
 import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException}
 
@@ -98,7 +96,12 @@ class SubExprEvaluationRuntime(cacheMaxEntries: Int) {
       val proxy = ExpressionProxy(expr, proxyExpressionCurrentId, this)
       proxyExpressionCurrentId += 1
 
-      proxyMap.putAll(e.map(_ -> proxy).toMap.asJava)
+      // We leverage `IdentityHashMap` so we compare expression keys by reference here.
+      // So for example if there are one group of common exprs like Seq(common expr 1,
+      // common expr2, ..., common expr n), we will insert into `proxyMap` some key/value
+      // pairs like Map(common expr 1 -> proxy(common expr 1), ...,
+      // common expr n -> proxy(common expr 1)).
+      e.map(proxyMap.put(_, proxy))
     }
 
     // Only adding proxy if we find subexpressions.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubExprEvaluationRuntimeSuite.scala

@@ -95,4 +95,26 @@ class SubExprEvaluationRuntimeSuite extends SparkFunSuite {
     })
     assert(proxys.isEmpty)
   }
+
+  test("SubExprEvaluationRuntime should wrap semantically equal exprs") {
+    val runtime = new SubExprEvaluationRuntime(1)
+
+    val one = Literal(1)
+    val two = Literal(2)
+    def mul: (Literal, Literal) => Expression =
+      (left: Literal, right: Literal) => Multiply(left, right)
+
+    val mul2_1 = Multiply(mul(one, two), mul(one, two))
+    val mul2_2 = Multiply(mul(one, two), mul(one, two))
+
+    val sqrt = Sqrt(mul2_1)
+    val sum = Add(mul2_2, sqrt)
+    val proxyExpressions = runtime.proxyExpressions(Seq(sum))
+    val proxys = proxyExpressions.flatMap(_.collect {
+      case p: ExpressionProxy => p
+    })
+    // ( (one * two) * (one * two) )
+    assert(proxys.size == 2)
+    assert(proxys.forall(_.child.semanticEquals(mul2_1)))
+  }
 }