[SPARK-6550][SQL] Add PreAnalyzer to keep logical plan consistent across DataFrame

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PreAnalyzer.scala

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+
+class PreAnalyzer(caseSensitive: Boolean = true,
+                  maxIterations: Int = 100) extends RuleExecutor[LogicalPlan] {
+
+  val resolver = if (caseSensitive) caseSensitiveResolution else caseInsensitiveResolution
+
+  val fixedPoint = FixedPoint(maxIterations)
+
+  lazy val batches: Seq[Batch] = Seq(
+    Batch("Resolution", fixedPoint, ResolveConflictingAttributes)
+  )
+
+  /**
+   * Handling the cases in which the attributes of nodes are conflicting
+   */
+  object ResolveConflictingAttributes extends Rule[LogicalPlan] {
+    def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
+      case p: LogicalPlan if !p.childrenResolved => p
+
+      // Special handling for cases when self-join introduce duplicate expression ids
+      case j @ Join(left, right, _, _) if left.outputSet.intersect(right.outputSet).nonEmpty =>
+        val conflictingAttributes = left.outputSet.intersect(right.outputSet)
+        logDebug(s"Conflicting attributes ${conflictingAttributes.mkString(",")} in $j")
+
+        val (oldRelation, newRelation) = right.collect {
+          // Handle base relations that might appear more than once.
+          case oldVersion: MultiInstanceRelation
+              if oldVersion.outputSet.intersect(conflictingAttributes).nonEmpty =>
+            val newVersion = oldVersion.newInstance()
+            (oldVersion, newVersion)
+
+          // Handle projects that create conflicting aliases.
+          case oldVersion @ Project(projectList, _)
+              if findAliases(projectList).intersect(conflictingAttributes).nonEmpty =>
+            (oldVersion, oldVersion.copy(projectList = newAliases(projectList)))
+
+          case oldVersion @ Aggregate(_, aggregateExpressions, _)
+              if findAliases(aggregateExpressions).intersect(conflictingAttributes).nonEmpty =>
+            (oldVersion, oldVersion.copy(aggregateExpressions = newAliases(aggregateExpressions)))
+        }.head // Only handle first case found, others will be fixed on the next pass.
+
+        val attributeRewrites = AttributeMap(oldRelation.output.zip(newRelation.output))
+        val newRight = right transformUp {
+          case r if r == oldRelation => newRelation
+        } transformUp {
+          case other => other transformExpressions {
+            case a: Attribute => attributeRewrites.get(a).getOrElse(a)
+          }
+        }
+        j.copy(right = newRight)
+    }
+
+    def newAliases(expressions: Seq[NamedExpression]): Seq[NamedExpression] = {
+      expressions.map {
+        case a: Alias => Alias(a.child, a.name)()
+        case other => other
+      }
+    }
+
+    def findAliases(projectList: Seq[NamedExpression]): AttributeSet = {
+      AttributeSet(projectList.collect { case a: Alias => a.toAttribute })
+    }
+  }
+}
+

sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala

@@ -113,6 +113,8 @@ class SQLContext(@transient val sparkContext: SparkContext)
   @transient
   protected[sql] lazy val functionRegistry: FunctionRegistry = new SimpleFunctionRegistry(true)
 
+  @transient
+  protected[sql] lazy val preAnalyzer: PreAnalyzer = new PreAnalyzer()
   @transient
   protected[sql] lazy val analyzer: Analyzer =
     new Analyzer(catalog, functionRegistry, caseSensitive = true) {
@@ -1104,9 +1106,10 @@ class SQLContext(@transient val sparkContext: SparkContext)
    * access to the intermediate phases of query execution for developers.
    */
   @DeveloperApi
-  protected[sql] class QueryExecution(val logical: LogicalPlan) {
+  protected[sql] class QueryExecution(val rawPlan: LogicalPlan) {
     def assertAnalyzed(): Unit = checkAnalysis(analyzed)
 
+    lazy val logical: LogicalPlan = preAnalyzer(rawPlan)
     lazy val analyzed: LogicalPlan = analyzer(logical)
     lazy val withCachedData: LogicalPlan = {
       assertAnalyzed()

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -113,6 +113,10 @@ class DataFrameSuite extends QueryTest {
     checkAnswer(
       df.as('x).join(df.as('y), $"x.str" === $"y.str").groupBy("x.str").count(),
       Row("1", 1) :: Row("2", 1) :: Row("3", 1) :: Nil)
+
+    checkAnswer(
+      df.as('x).join(df.as('y), $"x.str" === $"y.str").groupBy("y.str").count(),
+      Row("1", 1) :: Row("2", 1) :: Row("3", 1) :: Nil)
   }
 
   test("explode") {

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala

@@ -376,8 +376,11 @@ class HiveContext(sc: SparkContext) extends SQLContext(sc) {
   override protected[sql] val planner = hivePlanner
 
   /** Extends QueryExecution with hive specific features. */
-  protected[sql] class QueryExecution(logicalPlan: LogicalPlan)
-    extends super.QueryExecution(logicalPlan) {
+  protected[sql] class QueryExecution(rawPlan: LogicalPlan)
+    extends super.QueryExecution(rawPlan) {
+
+    lazy val logicalPlan: LogicalPlan = preAnalyzer(rawPlan)
+
     // Like what we do in runHive, makes sure the session represented by the
     // `sessionState` field is activated.
     if (SessionState.get() != sessionState) {