For review comment.

Author: viirya

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

@@ -16,7 +16,9 @@
  */
 package org.apache.spark.sql.catalyst.expressions
 
-import scala.collection.mutable
+import java.util.IdentityHashMap
+
+import scala.collection.JavaConverters._
 
 import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache}
 import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException}
@@ -32,6 +34,10 @@ import org.apache.spark.sql.types.DataType
  * intercepts expression evaluation and loads from the cache first.
  */
 class SubExprEvaluationRuntime(cacheMaxEntries: Int) {
+  // The id assigned to `ExpressionProxy`. `SubExprEvaluationRuntime` will use assigned ids of
+  // `ExpressionProxy` to decide the equality when loading from cache. `SubExprEvaluationRuntime`
+  // won't be use by multi-threads so we don't need to consider concurrency here.
+  private var proxyExpressionCurrentId = 0
 
   private[sql] val cache: LoadingCache[ExpressionProxy, ResultProxy] = CacheBuilder.newBuilder()
     .maximumSize(cacheMaxEntries)
@@ -68,8 +74,12 @@ class SubExprEvaluationRuntime(cacheMaxEntries: Int) {
    */
   private def replaceWithProxy(
       expr: Expression,
-      proxyMap: Map[Expression, ExpressionProxy]): Expression = {
-    proxyMap.getOrElse(expr, expr.mapChildren(replaceWithProxy(_, proxyMap)))
+      proxyMap: IdentityHashMap[Expression, ExpressionProxy]): Expression = {
+    if (proxyMap.containsKey(expr)) {
+      proxyMap.get(expr)
+    } else {
+      expr.mapChildren(replaceWithProxy(_, proxyMap))
+    }
   }
 
   /**
@@ -80,19 +90,20 @@ class SubExprEvaluationRuntime(cacheMaxEntries: Int) {
 
     expressions.foreach(equivalentExpressions.addExprTree(_))
 
-    val proxyMap = mutable.Map.empty[Expression, ExpressionProxy]
+    val proxyMap = new IdentityHashMap[Expression, ExpressionProxy]
 
     val commonExprs = equivalentExpressions.getAllEquivalentExprs.filter(_.size > 1)
     commonExprs.foreach { e =>
       val expr = e.head
-      val proxy = ExpressionProxy(expr, this)
+      val proxy = ExpressionProxy(expr, proxyExpressionCurrentId, this)
+      proxyExpressionCurrentId += 1
 
-      proxyMap ++= e.map(_ -> proxy).toMap
+      proxyMap.putAll(e.map(_ -> proxy).toMap.asJava)
     }
 
     // Only adding proxy if we find subexpressions.
-    if (proxyMap.nonEmpty) {
-      expressions.map(replaceWithProxy(_, proxyMap.toMap))
+    if (!proxyMap.isEmpty) {
+      expressions.map(replaceWithProxy(_, proxyMap))
     } else {
       expressions
     }
@@ -105,6 +116,7 @@ class SubExprEvaluationRuntime(cacheMaxEntries: Int) {
  */
 case class ExpressionProxy(
     child: Expression,
+    id: Int,
     runtime: SubExprEvaluationRuntime) extends Expression {
 
   final override def dataType: DataType = child.dataType
@@ -118,6 +130,13 @@ case class ExpressionProxy(
   def proxyEval(input: InternalRow = null): Any = child.eval(input)
 
   override def eval(input: InternalRow = null): Any = runtime.getEval(this)
+
+  override def equals(obj: Any): Boolean = obj match {
+    case other: ExpressionProxy => this.id == other.id
+    case _ => false
+  }
+
+  override def hashCode(): Int = this.id.hashCode()
 }
 
 /**

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

@@ -17,12 +17,13 @@
 package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.types.IntegerType
 
 class SubExprEvaluationRuntimeSuite extends SparkFunSuite {
 
   test("Evaluate ExpressionProxy should create cached result") {
     val runtime = new SubExprEvaluationRuntime(1)
-    val proxy = ExpressionProxy(Literal(1), runtime)
+    val proxy = ExpressionProxy(Literal(1), 0, runtime)
     assert(runtime.cache.size() == 0)
     proxy.eval()
     assert(runtime.cache.size() == 1)
@@ -33,31 +34,31 @@ class SubExprEvaluationRuntimeSuite extends SparkFunSuite {
     val runtime = new SubExprEvaluationRuntime(2)
     assert(runtime.cache.size() == 0)
 
-    val proxy1 = ExpressionProxy(Literal(1), runtime)
+    val proxy1 = ExpressionProxy(Literal(1), 0, runtime)
     proxy1.eval()
     assert(runtime.cache.size() == 1)
     assert(runtime.cache.get(proxy1) == ResultProxy(1))
 
-    val proxy2 = ExpressionProxy(Literal(2), runtime)
+    val proxy2 = ExpressionProxy(Literal(2), 1, runtime)
     proxy2.eval()
     assert(runtime.cache.size() == 2)
     assert(runtime.cache.get(proxy2) == ResultProxy(2))
 
-    val proxy3 = ExpressionProxy(Literal(3), runtime)
+    val proxy3 = ExpressionProxy(Literal(3), 2, runtime)
     proxy3.eval()
     assert(runtime.cache.size() == 2)
     assert(runtime.cache.get(proxy3) == ResultProxy(3))
   }
 
   test("setInput should empty cached result") {
     val runtime = new SubExprEvaluationRuntime(2)
-    val proxy1 = ExpressionProxy(Literal(1), runtime)
+    val proxy1 = ExpressionProxy(Literal(1), 0, runtime)
     assert(runtime.cache.size() == 0)
     proxy1.eval()
     assert(runtime.cache.size() == 1)
     assert(runtime.cache.get(proxy1) == ResultProxy(1))
 
-    val proxy2 = ExpressionProxy(Literal(2), runtime)
+    val proxy2 = ExpressionProxy(Literal(2), 1, runtime)
     proxy2.eval()
     assert(runtime.cache.size() == 2)
     assert(runtime.cache.get(proxy2) == ResultProxy(2))
@@ -83,8 +84,8 @@ class SubExprEvaluationRuntimeSuite extends SparkFunSuite {
     })
     // ( (one * two) * (one * two) )
     assert(proxys.size == 2)
-    val expected = ExpressionProxy(mul2, runtime)
-    assert(proxys.head == expected)
+    val expected = ExpressionProxy(mul2, 0, runtime)
+    assert(proxys.forall(_ == expected))
   }
 
   test("ExpressionProxy won't be on non deterministic") {