[SPARK-28344][SQL] detect ambiguous self-join and fail the query

This is an alternative solution of #24442 . It fails the query if ambiguous self join is detected, instead of trying to disambiguate it. The problem is that, it's hard to come up with a reasonable rule to disambiguate, the rule proposed by #24442 is mostly a heuristic.

This is a long-standing bug and I've seen many people complaining about it in JIRA/dev list.

A typical example:
```
val df1 = …
val df2 = df1.filter(...)
df1.join(df2, df1("a") > df2("a")) // returns empty result
```
The root cause is, `Dataset.apply` is so powerful that users think it returns a column reference which can point to the column of the Dataset at anywhere. This is not true in many cases. `Dataset.apply` returns an `AttributeReference` . Different Datasets may share the same `AttributeReference`. In the example above, `df2` adds a Filter operator above the logical plan of `df1`, and the Filter operator reserves the output `AttributeReference` of its child. This means, `df1("a")` is exactly the same as `df2("a")`, and `df1("a") > df2("a")` always evaluates to false.

We can reuse the infra in #24442 :
1. each Dataset has a globally unique id.
2. the `AttributeReference` returned by `Dataset.apply` carries the ID and column position(e.g. 3rd column of the Dataset) via metadata.
3. the logical plan of a `Dataset` carries the ID via `TreeNodeTag`

When self-join happens, the analyzer asks the right side plan of join to re-generate output attributes with new exprIds. Based on it, a simple rule to detect ambiguous self join is:
1. find all column references (i.e. `AttributeReference`s with Dataset ID and col position) in the root node of a query plan.
2. for each column reference, traverse the query plan tree, find a sub-plan that carries Dataset ID and the ID is the same as the one in the column reference.
3. get the corresponding output attribute of the sub-plan by the col position in the column reference.
4. if the corresponding output attribute has a different exprID than the column reference, then it means this sub-plan is on the right side of a self-join and has regenerated its output attributes. This is an ambiguous self join because the column reference points to a table being self-joined.

existing tests and new test cases

Closes #25107 from cloud-fan/new-self-join.

Authored-by: Wenchen Fan <wenchen@databricks.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: cloud-fan

docs/sql-migration-guide-upgrade.md

@@ -16,11 +16,13 @@ displayTitle: Spark SQL Upgrading Guide
 
  - Since Spark 2.4.5, `TRUNCATE TABLE` command tries to set back original permission and ACLs during re-creating the table/partition paths. To restore the behaviour of earlier versions, set `spark.sql.truncateTable.ignorePermissionAcl.enabled` to `true`.
 
-  - Since Spark 2.4.5, `spark.sql.legacy.mssqlserver.numericMapping.enabled` configuration is added in order to support the legacy MsSQLServer dialect mapping behavior using IntegerType and DoubleType for SMALLINT and REAL JDBC types, respectively. To restore the behaviour of 2.4.3 and earlier versions, set `spark.sql.legacy.mssqlserver.numericMapping.enabled` to `true`.
+ - Since Spark 2.4.5, `spark.sql.legacy.mssqlserver.numericMapping.enabled` configuration is added in order to support the legacy MsSQLServer dialect mapping behavior using IntegerType and DoubleType for SMALLINT and REAL JDBC types, respectively. To restore the behaviour of 2.4.3 and earlier versions, set `spark.sql.legacy.mssqlserver.numericMapping.enabled` to `true`.
+
+ - Since Spark 2.4.5, Dataset query fails if it contains ambiguous column reference that is caused by self join. A typical example: `val df1 = ...; val df2 = df1.filter(...);`, then `df1.join(df2, df1("a") > df2("a"))` returns an empty result which is quite confusing. This is because Spark cannot resolve Dataset column references that point to tables being self joined, and `df1("a")` is exactly the same as `df2("a")` in Spark. To restore the behavior before Spark 3.0, you can set `spark.sql.analyzer.failAmbiguousSelfJoin` to `false`.
 
 ## Upgrading from Spark SQL 2.4.3 to 2.4.4
 
-    - Since Spark 2.4.4, according to [MsSqlServer Guide](https://docs.microsoft.com/en-us/sql/connect/jdbc/using-basic-data-types?view=sql-server-2017), MsSQLServer JDBC Dialect uses ShortType and FloatType for SMALLINT and REAL, respectively. Previously, IntegerType and DoubleType is used.
+ - Since Spark 2.4.4, according to [MsSqlServer Guide](https://docs.microsoft.com/en-us/sql/connect/jdbc/using-basic-data-types?view=sql-server-2017), MsSQLServer JDBC Dialect uses ShortType and FloatType for SMALLINT and REAL, respectively. Previously, IntegerType and DoubleType is used.
 
 ## Upgrading from Spark SQL 2.4 to 2.4.1
 

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

@@ -334,7 +334,8 @@ class Analyzer(
         gid: Expression): Expression = {
       expr transform {
         case e: GroupingID =>
-          if (e.groupByExprs.isEmpty || e.groupByExprs == groupByExprs) {
+          if (e.groupByExprs.isEmpty ||
+              e.groupByExprs.map(_.canonicalized) == groupByExprs.map(_.canonicalized)) {
             Alias(gid, toPrettySQL(e))()
           } else {
             throw new AnalysisException(
@@ -936,6 +937,8 @@ class Analyzer(
       // To resolve duplicate expression IDs for Join and Intersect
       case j @ Join(left, right, _, _) if !j.duplicateResolved =>
         j.copy(right = dedupRight(left, right))
+      // intersect/except will be rewritten to join at the begininng of optimizer. Here we need to
+      // deduplicate the right side plan, so that we won't produce an invalid self-join later.
       case i @ Intersect(left, right, _) if !i.duplicateResolved =>
         i.copy(right = dedupRight(left, right))
       case e @ Except(left, right, _) if !e.duplicateResolved =>

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -726,6 +726,13 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val FAIL_AMBIGUOUS_SELF_JOIN =
+    buildConf("spark.sql.analyzer.failAmbiguousSelfJoin")
+      .doc("When true, fail the Dataset query if it contains ambiguous self-join.")
+      .internal()
+      .booleanConf
+      .createWithDefault(true)
+
   // Whether to retain group by columns or not in GroupedData.agg.
   val DATAFRAME_RETAIN_GROUP_COLUMNS = buildConf("spark.sql.retainGroupColumns")
     .internal()

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

@@ -48,6 +48,15 @@ private[sql] object Column {
       case expr => toPrettySQL(expr)
     }
   }
+
+  private[sql] def stripColumnReferenceMetadata(a: AttributeReference): AttributeReference = {
+    val metadataWithoutId = new MetadataBuilder()
+      .withMetadata(a.metadata)
+      .remove(Dataset.DATASET_ID_KEY)
+      .remove(Dataset.COL_POS_KEY)
+      .build()
+    a.withMetadata(metadataWithoutId)
+  }
 }
 
 /**
@@ -141,11 +150,15 @@ class Column(val expr: Expression) extends Logging {
   override def toString: String = toPrettySQL(expr)
 
   override def equals(that: Any): Boolean = that match {
-    case that: Column => that.expr.equals(this.expr)
+    case that: Column => that.normalizedExpr() == this.normalizedExpr()
     case _ => false
   }
 
-  override def hashCode: Int = this.expr.hashCode()
+  override def hashCode: Int = this.normalizedExpr().hashCode()
+
+  private def normalizedExpr(): Expression = expr transform {
+    case a: AttributeReference => Column.stripColumnReferenceMetadata(a)
+  }
 
   /** Creates a column based on the given expression. */
   private def withExpr(newExpr: Expression): Column = new Column(newExpr)
@@ -1023,7 +1036,7 @@ class Column(val expr: Expression) extends Logging {
    * @since 2.0.0
    */
   def name(alias: String): Column = withExpr {
-    expr match {
+    normalizedExpr() match {
       case ne: NamedExpression => Alias(expr, alias)(explicitMetadata = Some(ne.metadata))
       case other => Alias(other, alias)()
     }

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

@@ -45,12 +45,14 @@ import org.apache.spark.sql.catalyst.parser.{ParseException, ParserUtils}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, PartitioningCollection}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.arrow.{ArrowBatchStreamWriter, ArrowConverters}
 import org.apache.spark.sql.execution.command._
 import org.apache.spark.sql.execution.datasources.LogicalRelation
 import org.apache.spark.sql.execution.python.EvaluatePython
 import org.apache.spark.sql.execution.stat.StatFunctions
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.DataStreamWriter
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.util.SchemaUtils
@@ -60,6 +62,11 @@ import org.apache.spark.unsafe.types.CalendarInterval
 import org.apache.spark.util.Utils
 
 private[sql] object Dataset {
+  val curId = new java.util.concurrent.atomic.AtomicLong()
+  val DATASET_ID_KEY = "__dataset_id"
+  val COL_POS_KEY = "__col_position"
+  val DATASET_ID_TAG = TreeNodeTag[Long]("dataset_id")
+
   def apply[T: Encoder](sparkSession: SparkSession, logicalPlan: LogicalPlan): Dataset[T] = {
     val dataset = new Dataset(sparkSession, logicalPlan, implicitly[Encoder[T]])
     // Eagerly bind the encoder so we verify that the encoder matches the underlying
@@ -173,6 +180,9 @@ class Dataset[T] private[sql](
     encoder: Encoder[T])
   extends Serializable {
 
+  // A globally unique id of this Dataset.
+  private val id = Dataset.curId.getAndIncrement()
+
   queryExecution.assertAnalyzed()
 
   // Note for Spark contributors: if adding or updating any action in `Dataset`, please make sure
@@ -189,14 +199,18 @@ class Dataset[T] private[sql](
   @transient private[sql] val logicalPlan: LogicalPlan = {
     // For various commands (like DDL) and queries with side effects, we force query execution
     // to happen right away to let these side effects take place eagerly.
-    queryExecution.analyzed match {
+    val plan = queryExecution.analyzed match {
       case c: Command =>
         LocalRelation(c.output, withAction("command", queryExecution)(_.executeCollect()))
       case u @ Union(children) if children.forall(_.isInstanceOf[Command]) =>
         LocalRelation(u.output, withAction("command", queryExecution)(_.executeCollect()))
       case _ =>
         queryExecution.analyzed
     }
+    if (sparkSession.sessionState.conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN)) {
+      plan.setTagValue(Dataset.DATASET_ID_TAG, id)
+    }
+    plan
   }
 
   /**
@@ -1271,11 +1285,29 @@ class Dataset[T] private[sql](
       if (sqlContext.conf.supportQuotedRegexColumnName) {
         colRegex(colName)
       } else {
-        val expr = resolve(colName)
-        Column(expr)
+        Column(addDataFrameIdToCol(resolve(colName)))
       }
   }
 
+  // Attach the dataset id and column position to the column reference, so that we can detect
+  // ambiguous self-join correctly. See the rule `DetectAmbiguousSelfJoin`.
+  // This must be called before we return a `Column` that contains `AttributeReference`.
+  // Note that, the metadata added here are only avaiable in the analyzer, as the analyzer rule
+  // `DetectAmbiguousSelfJoin` will remove it.
+  private def addDataFrameIdToCol(expr: NamedExpression): NamedExpression = {
+    val newExpr = expr transform {
+      case a: AttributeReference
+        if sparkSession.sessionState.conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN) =>
+        val metadata = new MetadataBuilder()
+          .withMetadata(a.metadata)
+          .putLong(Dataset.DATASET_ID_KEY, id)
+          .putLong(Dataset.COL_POS_KEY, logicalPlan.output.indexWhere(a.semanticEquals))
+          .build()
+        a.withMetadata(metadata)
+    }
+    newExpr.asInstanceOf[NamedExpression]
+  }
+
   /**
    * Selects column based on the column name specified as a regex and returns it as [[Column]].
    * @group untypedrel
@@ -1289,7 +1321,7 @@ class Dataset[T] private[sql](
       case ParserUtils.qualifiedEscapedIdentifier(nameParts, columnNameRegex) =>
         Column(UnresolvedRegex(columnNameRegex, Some(nameParts), caseSensitive))
       case _ =>
-        Column(resolve(colName))
+        Column(addDataFrameIdToCol(resolve(colName)))
     }
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/analysis/DetectAmbiguousSelfJoin.scala

@@ -0,0 +1,162 @@
+/*
+ * 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.execution.analysis
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.{AnalysisException, Column, Dataset}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Cast, Equality, Expression, ExprId}
+import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Detects ambiguous self-joins, so that we can fail the query instead of returning confusing
+ * results.
+ *
+ * Dataset column reference is simply an [[AttributeReference]] that is returned by `Dataset#col`.
+ * Most of time we don't need to do anything special, as [[AttributeReference]] can point to
+ * the column precisely. However, in case of self-join, the analyzer generates
+ * [[AttributeReference]] with new expr IDs for the right side plan of the join. If the Dataset
+ * column reference points to a column in the right side plan of a self-join, users will get
+ * unexpected result because the column reference can't match the newly generated
+ * [[AttributeReference]].
+ *
+ * Note that, this rule removes all the Dataset id related metadata from `AttributeReference`, so
+ * that they don't exist after analyzer.
+ */
+class DetectAmbiguousSelfJoin(conf: SQLConf) extends Rule[LogicalPlan] {
+
+  // Dataset column reference is an `AttributeReference` with 2 special metadata.
+  private def isColumnReference(a: AttributeReference): Boolean = {
+    a.metadata.contains(Dataset.DATASET_ID_KEY) && a.metadata.contains(Dataset.COL_POS_KEY)
+  }
+
+  private case class ColumnReference(datasetId: Long, colPos: Int, exprId: ExprId)
+
+  private def toColumnReference(a: AttributeReference): ColumnReference = {
+    ColumnReference(
+      a.metadata.getLong(Dataset.DATASET_ID_KEY),
+      a.metadata.getLong(Dataset.COL_POS_KEY).toInt,
+      a.exprId)
+  }
+
+  object LogicalPlanWithDatasetId {
+    def unapply(p: LogicalPlan): Option[(LogicalPlan, Long)] = {
+      p.getTagValue(Dataset.DATASET_ID_TAG).map(id => p -> id)
+    }
+  }
+
+  object AttrWithCast {
+    def unapply(expr: Expression): Option[AttributeReference] = expr match {
+      case Cast(child, _, _) => unapply(child)
+      case a: AttributeReference => Some(a)
+      case _ => None
+    }
+  }
+
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    if (!conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN)) return plan
+
+    // We always remove the special metadata from `AttributeReference` at the end of this rule, so
+    // Dataset column reference only exists in the root node via Dataset transformations like
+    // `Dataset#select`.
+    val colRefAttrs = plan.expressions.flatMap(_.collect {
+      case a: AttributeReference if isColumnReference(a) => a
+    })
+
+    if (colRefAttrs.nonEmpty) {
+      val colRefs = colRefAttrs.map(toColumnReference).distinct
+      val ambiguousColRefs = mutable.HashSet.empty[ColumnReference]
+      val dsIdSet = colRefs.map(_.datasetId).toSet
+
+      plan.foreach {
+        case LogicalPlanWithDatasetId(p, id) if dsIdSet.contains(id) =>
+          colRefs.foreach { ref =>
+            if (id == ref.datasetId) {
+              if (ref.colPos < 0 || ref.colPos >= p.output.length) {
+                throw new IllegalStateException("[BUG] Hit an invalid Dataset column reference: " +
+                  s"$ref. Please open a JIRA ticket to report it.")
+              } else {
+                // When self-join happens, the analyzer asks the right side plan to generate
+                // attributes with new exprIds. If a plan of a Dataset outputs an attribute which
+                // is referred by a column reference, and this attribute has different exprId than
+                // the attribute of column reference, then the column reference is ambiguous, as it
+                // refers to a column that gets regenerated by self-join.
+                val actualAttr = p.output(ref.colPos).asInstanceOf[AttributeReference]
+                if (actualAttr.exprId != ref.exprId) {
+                  ambiguousColRefs += ref
+                }
+              }
+            }
+          }
+
+        case _ =>
+      }
+
+      val ambiguousAttrs: Seq[AttributeReference] = plan match {
+        case Join(
+            LogicalPlanWithDatasetId(_, leftId),
+            LogicalPlanWithDatasetId(_, rightId),
+            _, condition, _) =>
+          // If we are dealing with root join node, we need to take care of SPARK-6231:
+          //  1. We can de-ambiguous `df("col") === df("col")` in the join condition.
+          //  2. There is no ambiguity in direct self join like
+          //     `df.join(df, df("col") === 1)`, because it doesn't matter which side the
+          //     column comes from.
+          def getAmbiguousAttrs(expr: Expression): Seq[AttributeReference] = expr match {
+            case Equality(AttrWithCast(a), AttrWithCast(b)) if a.sameRef(b) =>
+              Nil
+            case Equality(AttrWithCast(a), b) if leftId == rightId && b.foldable =>
+              Nil
+            case Equality(a, AttrWithCast(b)) if leftId == rightId && a.foldable =>
+              Nil
+            case a: AttributeReference =>
+              if (isColumnReference(a)) {
+                val colRef = toColumnReference(a)
+                if (ambiguousColRefs.contains(colRef)) Seq(a) else Nil
+              } else {
+                Nil
+              }
+            case _ => expr.children.flatMap(getAmbiguousAttrs)
+          }
+          condition.toSeq.flatMap(getAmbiguousAttrs)
+
+        case _ => ambiguousColRefs.toSeq.map { ref =>
+          colRefAttrs.find(attr => toColumnReference(attr) == ref).get
+        }
+      }
+
+      if (ambiguousAttrs.nonEmpty) {
+        throw new AnalysisException(s"Column ${ambiguousAttrs.mkString(", ")} are ambiguous. " +
+          "It's probably because you joined several Datasets together, and some of these " +
+          "Datasets are the same. This column points to one of the Datasets but Spark is unable " +
+          "to figure out which one. Please alias the Datasets with different names via " +
+          "`Dataset.as` before joining them, and specify the column using qualified name, e.g. " +
+          """`df.as("a").join(df.as("b"), $"a.id" > $"b.id")`. You can also set """ +
+          s"${SQLConf.FAIL_AMBIGUOUS_SELF_JOIN.key} to false to disable this check.")
+      }
+    }
+
+    plan.transformExpressions {
+      case a: AttributeReference if isColumnReference(a) =>
+        // Remove the special metadata from this `AttributeReference`, as the detection is done.
+        Column.stripColumnReferenceMetadata(a)
+    }
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala

@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.parser.ParserInterface
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.{QueryExecution, SparkOptimizer, SparkPlanner, SparkSqlParser}
+import org.apache.spark.sql.execution.analysis.DetectAmbiguousSelfJoin
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.streaming.StreamingQueryManager
 import org.apache.spark.sql.util.ExecutionListenerManager
@@ -161,7 +162,8 @@ abstract class BaseSessionStateBuilder(
         customResolutionRules
 
     override val postHocResolutionRules: Seq[Rule[LogicalPlan]] =
-      PreprocessTableCreation(session) +:
+      new DetectAmbiguousSelfJoin(conf) +:
+        PreprocessTableCreation(session) +:
         PreprocessTableInsertion(conf) +:
         DataSourceAnalysis(conf) +:
         customPostHocResolutionRules

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

@@ -167,6 +167,21 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
         Row(null, null, 1, 1, 3) :: Nil
     )
 
+    // use column reference in `grouping_id` instead of column name
+    checkAnswer(
+      courseSales.cube("course", "year")
+        .agg(grouping_id(courseSales("course"), courseSales("year"))),
+      Row("Java", 2012, 0) ::
+        Row("Java", 2013, 0) ::
+        Row("Java", null, 1) ::
+        Row("dotNET", 2012, 0) ::
+        Row("dotNET", 2013, 0) ::
+        Row("dotNET", null, 1) ::
+        Row(null, 2012, 2) ::
+        Row(null, 2013, 2) ::
+        Row(null, null, 3) :: Nil
+    )
+
     intercept[AnalysisException] {
       courseSales.groupBy().agg(grouping("course")).explain()
     }