From f3dd17a5564ce5098aab112dda28fa02cd802af6 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 29 Feb 2024 15:20:02 -0800
Subject: [PATCH 01/22] SPARK-47217. bug fix for exception thrown in reused
 dataframes involving joins once the plan is de-duplicated. The fix involves
 using Dataset ID associated with the plans & attributes to attempt correct
 resolution

---
 .../analysis/ColumnResolutionHelper.scala     | 42 ++++++++++
 .../sql/catalyst/analysis/unresolved.scala    | 41 ++++++++++
 .../catalyst/plans/logical/LogicalPlan.scala  |  4 +-
 .../scala/org/apache/spark/sql/Dataset.scala  | 37 ++++++++-
 .../apache/spark/sql/JavaDataFrameSuite.java  | 76 +++++++++++++++++++
 .../spark/sql/DataFrameAsOfJoinSuite.scala    | 20 +++++
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 10 +++
 7 files changed, 226 insertions(+), 4 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index 8ea50e2ceb659..22361b0e73c77 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -134,6 +134,7 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       expr: Expression,
       resolveColumnByName: Seq[String] => Option[Expression],
       getAttrCandidates: () => Seq[Attribute],
+      resolveOnDatasetId: (Long, String) => Option[NamedExpression],
       throws: Boolean,
       includeLastResort: Boolean): Expression = {
     def innerResolve(e: Expression, isTopLevel: Boolean): Expression = withOrigin(e.origin) {
@@ -156,6 +157,9 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
           }
           matched(ordinal)
 
+        case u @ UnresolvedAttributeWithTag(attr, id) =>
+          resolveOnDatasetId(id, attr.name).getOrElse(attr)
+
         case u @ UnresolvedAttribute(nameParts) =>
           val result = withPosition(u) {
             resolveColumnByName(nameParts).orElse(resolveLiteralFunction(nameParts)).map {
@@ -452,6 +456,7 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
         plan.resolve(nameParts, conf.resolver)
       },
       getAttrCandidates = () => plan.output,
+      resolveOnDatasetId = (_, _) => None,
       throws = throws,
       includeLastResort = includeLastResort)
   }
@@ -477,6 +482,43 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
         assert(q.children.length == 1)
         q.children.head.output
       },
+
+      resolveOnDatasetId = (datasetid: Long, name: String) => {
+        def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[LogicalPlan] = {
+          var currentLp = lp
+          while(currentLp.children.size < 2) {
+            if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
+              return Option(currentLp)
+            } else {
+              if (currentLp.children.size == 1) {
+                currentLp = currentLp.children.head
+              } else {
+                // leaf node
+                return None
+              }
+            }
+          }
+          None
+        }
+
+        val binaryNodeOpt = q.collectFirst {
+          case bn: BinaryNode => bn
+        }
+
+        val resolveOnAttribs = binaryNodeOpt match {
+          case Some(bn) =>
+            val leftDefOpt = findUnaryNodeMatchingTagId(bn.left)
+            val rightDefOpt = findUnaryNodeMatchingTagId(bn.right)
+            (leftDefOpt, rightDefOpt) match {
+              case (None, Some(lp)) => lp.output
+              case (Some(lp), None) => lp.output
+              case _ => q.children.head.output
+            }
+
+          case _ => q.children.head.output
+        }
+        AttributeSeq.fromNormalOutput(resolveOnAttribs).resolve(Seq(name), conf.resolver)
+      },
       throws = true,
       includeLastResort = includeLastResort)
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index 7a3cc4bc8e83e..397351e0c1fdd 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -268,6 +268,47 @@ case class UnresolvedAttribute(nameParts: Seq[String]) extends Attribute with Un
   }
 }
 
+case class UnresolvedAttributeWithTag(attribute: Attribute, datasetId: Long) extends Attribute with
+  Unevaluable {
+  def name: String = attribute.name
+
+  override def exprId: ExprId = throw new UnresolvedException("exprId")
+
+  override def dataType: DataType = throw new UnresolvedException("dataType")
+
+  override def nullable: Boolean = throw new UnresolvedException("nullable")
+
+  override def qualifier: Seq[String] = throw new UnresolvedException("qualifier")
+
+  override lazy val resolved = false
+
+  override def newInstance(): UnresolvedAttributeWithTag = this
+
+  override def withNullability(newNullability: Boolean): UnresolvedAttributeWithTag = this
+
+  override def withQualifier(newQualifier: Seq[String]): UnresolvedAttributeWithTag = this
+
+  override def withName(newName: String): UnresolvedAttributeWithTag = this
+
+  override def withMetadata(newMetadata: Metadata): Attribute = this
+
+  override def withExprId(newExprId: ExprId): UnresolvedAttributeWithTag = this
+
+  override def withDataType(newType: DataType): Attribute = this
+
+  final override val nodePatterns: Seq[TreePattern] = Seq(UNRESOLVED_ATTRIBUTE)
+
+  override def toString: String = s"'$name"
+
+  override def sql: String = attribute.sql
+
+  /**
+   * Returns true if this matches the token. This requires the attribute to only have one part in
+   * its name and that matches the given token in a case insensitive way.
+   */
+  def equalsIgnoreCase(token: String): Boolean = token.equalsIgnoreCase(attribute.name)
+}
+
 object UnresolvedAttribute extends AttributeNameParser {
   /**
    * Creates an [[UnresolvedAttribute]], parsing segments separated by dots ('.').
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
index e1121d1f9026e..a9b130c981ac0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.plans.logical
 
+import scala.collection.mutable
+
 import org.apache.spark.SparkUnsupportedOperationException
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
@@ -30,7 +32,6 @@ import org.apache.spark.sql.catalyst.util.MetadataColumnHelper
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
 import org.apache.spark.sql.types.{MapType, StructType}
 
-
 abstract class LogicalPlan
   extends QueryPlan[LogicalPlan]
   with AnalysisHelper
@@ -199,6 +200,7 @@ object LogicalPlan {
   //    to the old code path.
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
+  private[spark] val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 189be1d6a30d2..b767cc01f341d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -47,7 +47,7 @@ import org.apache.spark.sql.catalyst.json.{JacksonGenerator, JSONOptions}
 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.trees.{TreeNodeTag, TreePattern}
+import org.apache.spark.sql.catalyst.trees.TreePattern
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, IntervalUtils}
 import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId
@@ -73,7 +73,7 @@ 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[HashSet[Long]]("dataset_id")
+  val DATASET_ID_TAG = LogicalPlan.DATASET_ID_TAG
 
   def apply[T: Encoder](sparkSession: SparkSession, logicalPlan: LogicalPlan): Dataset[T] = {
     val dataset = new Dataset(sparkSession, logicalPlan, implicitly[Encoder[T]])
@@ -1308,12 +1308,20 @@ class Dataset[T] private[sql](
       case a: AttributeReference if logicalPlan.outputSet.contains(a) =>
         val index = logicalPlan.output.indexWhere(_.exprId == a.exprId)
         joined.left.output(index)
+
+      case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
+        UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
     }
+
     val rightAsOfExpr = rightAsOf.expr.transformUp {
       case a: AttributeReference if other.logicalPlan.outputSet.contains(a) =>
         val index = other.logicalPlan.output.indexWhere(_.exprId == a.exprId)
         joined.right.output(index)
+
+      case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
+        UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
     }
+
     withPlan {
       AsOfJoin(
         joined.left, joined.right,
@@ -1576,7 +1584,30 @@ class Dataset[T] private[sql](
 
       case other => other
     }
-    Project(untypedCols.map(_.named), logicalPlan)
+    val namedExprs = untypedCols.map(_.named)
+    val inputSet = logicalPlan.outputSet
+    val rectifiedNamedExprs = namedExprs.map(ne => ne match {
+
+      case al: Alias if !al.references.subsetOf(inputSet) &&
+        al.nonInheritableMetadataKeys.contains(Dataset.DATASET_ID_KEY) =>
+        val unresolvedExpr = al.child.transformUp {
+          case attr: AttributeReference if !inputSet.contains(attr) =>
+            UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
+        }
+        val newAl = al.copy(child = unresolvedExpr, name = al.name)(exprId = al.exprId,
+          qualifier = al.qualifier, explicitMetadata = al.explicitMetadata,
+          nonInheritableMetadataKeys = al.nonInheritableMetadataKeys)
+        newAl.copyTagsFrom(al)
+        newAl
+
+      case attr: Attribute if !inputSet.contains(attr) &&
+        attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
+        UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
+
+      case _ => ne
+
+    })
+    Project(rectifiedNamedExprs, logicalPlan)
   }
 
   /**
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 26a19cbed1b91..f325e3c1a80ec 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -24,6 +24,7 @@
 import java.math.BigInteger;
 import java.math.BigDecimal;
 
+import org.apache.spark.sql.catalyst.plans.logical.Join;
 import scala.collection.Seq;
 import scala.jdk.javaapi.CollectionConverters;
 
@@ -31,6 +32,11 @@
 import com.google.common.primitives.Ints;
 import org.junit.jupiter.api.*;
 
+import org.apache.spark.sql.catalyst.expressions.Alias;
+import org.apache.spark.sql.catalyst.expressions.AttributeReference;
+import org.apache.spark.sql.catalyst.expressions.AttributeSet;
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
+import org.apache.spark.sql.catalyst.plans.logical.Project;
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.Column;
@@ -540,4 +546,74 @@ public void testUDF() {
       .map(row -> row.get(0).toString() + row.getString(1)).toArray(String[]::new);
     Assertions.assertArrayEquals(expected, result);
   }
+
+  @Test
+  public void testDedupBehaviourOnProjection_SPARK_47217() {
+    //Create table1 DF
+    List<Row> table1Data = Arrays.asList(
+      RowFactory.create(1, 2, 3),
+      RowFactory.create(1, 2, 3));
+    StructType table1Scema = new StructType()
+      .add("col11", DataTypes.IntegerType)
+      .add("col12", DataTypes.IntegerType)
+      .add("col13", DataTypes.IntegerType);
+
+    Dataset<Row> table1 = spark.createDataFrame(table1Data, table1Scema);
+
+    //Create table2 DataFrame
+    List<Row> table2Data = Arrays.asList(RowFactory.create(1, 2, 3));
+    StructType table2Schema = new StructType()
+      .add("col21", DataTypes.IntegerType)
+      .add("col22", DataTypes.IntegerType)
+      .add("col23", DataTypes.IntegerType);
+
+    Dataset<Row> table2 = spark.createDataFrame(table2Data, table2Schema);
+
+    //Create table 3 DataFrame
+    List<Row> table3Data = Arrays.asList(RowFactory.create(1, 2, 3), RowFactory.create(1, 2, 3));
+    StructType table3Schema = new StructType().
+      add("col31", DataTypes.IntegerType).
+      add("col32", DataTypes.IntegerType).
+      add("col33", DataTypes.IntegerType);
+
+    Dataset<Row> table3 = spark.createDataFrame(table3Data, table3Schema);
+
+    //Perform left outer join for table2
+    Dataset<Row> srcDf = table1.join(
+        table2,
+        table1.col("col11").equalTo(table2.col("col21")),
+        "left_outer").select(
+        table1.col("col11"),
+        table1.col("col12"),
+        table1.col("col13"),
+        table2.col("col22"));
+
+    //Perform leftouter join for exchange table2(firstjoin)
+    srcDf = srcDf.join(
+            broadcast(table3),
+            srcDf.col("col12").equalTo(
+                table3.col("col32")), "left_outer").
+        select(srcDf.col("col11"), srcDf.col("col12"),
+            srcDf.col("col13"),
+            table3.col("col33").as("col33_1"));
+
+    //Perform left outer joinfor exchangeRateTable1 again(secondjoin)
+    Dataset<Row> temp = srcDf.join(broadcast(table3),
+        srcDf.col("col11").equalTo(
+            table3.col("col31")), "left_outer");
+
+    srcDf = temp.select(srcDf.col("col11"), srcDf.col("col12"),
+        srcDf.col("col13"), srcDf.col("col33_1"),
+        table3.col("col33").as("col33_2"));
+
+    // verify optimized plan creation ok
+    srcDf.queryExecution().optimizedPlan();
+    LogicalPlan lp = srcDf.queryExecution().analyzed();
+    // verify attribute ref resolution is correct, i.e it resolves to the right leg of join
+    AttributeReference refToCheck =
+        (AttributeReference) ((Alias)((Project)lp).projectList().last()).child();
+    AttributeSet compareSet =
+        ((Join) ((Project) lp).child()).right().outputSet();
+    Assertions.assertTrue(compareSet.contains(refToCheck));
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
index 280eb095dc753..3db8d374c4bbc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql
 
 import scala.jdk.CollectionConverters._
 
+import org.apache.spark.sql.catalyst.plans.logical.AsOfJoin
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSparkSession
@@ -173,4 +174,23 @@ class DataFrameAsOfJoinSuite extends QueryTest
       )
     )
   }
+
+  test("SPARK_47217: Dedup of relations can impact projected columns resolution -1") {
+    val (df1, df2) = prepareForAsOfJoin()
+    val join1 = df1.join(df2, df1.col("a") === df2.col("a")).select(df2.col("a"), df1.col("b"),
+      df2.col("b"), df1.col("a").as("aa"))
+
+    // In stock spark this would throw ambiguous column exception, even though it is not ambiguous
+    val asOfjoin2 = join1.joinAsOf(
+      df1, df1.col("a"), join1.col("a"), usingColumns = Seq.empty,
+      joinType = "left", tolerance = null, allowExactMatches = false, direction = "nearest")
+
+    asOfjoin2.queryExecution.assertAnalyzed()
+
+    val testDf = asOfjoin2.select(df1.col("a"))
+    val analyzed = testDf.queryExecution.analyzed
+    val attributeRefToCheck = analyzed.output.head
+    assert(analyzed.children(0).asInstanceOf[AsOfJoin].right.outputSet.
+      contains(attributeRefToCheck))
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index c777d2207584d..83d289fcd2073 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -488,4 +488,14 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
         context = ExpectedContext(fragment = "$", getCurrentClassCallSitePattern))
     }
   }
+
+  test("SPARK_47217: Dedup of relations can impact projected columns resolution") {
+    val df = Seq((1, 2)).toDF("a", "b")
+    val df2 = df.select(df("a").as("aa"), df("b").as("bb"))
+    val df3 = df2.join(df, df2("bb") === df("b")).select(df2("aa"), df("a"))
+
+    checkAnswer(
+      df3,
+      Row(1, 1) :: Nil)
+  }
 }

From c29366f489a934eb6aff8b029244fda5cc77daef Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 29 Feb 2024 16:52:22 -0800
Subject: [PATCH 02/22] SPARK-47217. fix test failures

---
 sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index b767cc01f341d..7a3eae7144b37 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1591,7 +1591,8 @@ class Dataset[T] private[sql](
       case al: Alias if !al.references.subsetOf(inputSet) &&
         al.nonInheritableMetadataKeys.contains(Dataset.DATASET_ID_KEY) =>
         val unresolvedExpr = al.child.transformUp {
-          case attr: AttributeReference if !inputSet.contains(attr) =>
+          case attr: AttributeReference if !inputSet.contains(attr) &&
+            attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
             UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
         }
         val newAl = al.copy(child = unresolvedExpr, name = al.name)(exprId = al.exprId,

From 31d66c2540423aa852d444c14e88cbfcfaaedc8c Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 29 Feb 2024 19:11:17 -0800
Subject: [PATCH 03/22] SPARK-47217. fix style format issue

---
 .../apache/spark/sql/JavaDataFrameSuite.java    | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index f325e3c1a80ec..09ccbd508e454 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -571,10 +571,10 @@ public void testDedupBehaviourOnProjection_SPARK_47217() {
 
     //Create table 3 DataFrame
     List<Row> table3Data = Arrays.asList(RowFactory.create(1, 2, 3), RowFactory.create(1, 2, 3));
-    StructType table3Schema = new StructType().
-      add("col31", DataTypes.IntegerType).
-      add("col32", DataTypes.IntegerType).
-      add("col33", DataTypes.IntegerType);
+    StructType table3Schema = new StructType()
+      .add("col31", DataTypes.IntegerType)
+      .add("col32", DataTypes.IntegerType)
+      .add("col33", DataTypes.IntegerType);
 
     Dataset<Row> table3 = spark.createDataFrame(table3Data, table3Schema);
 
@@ -591,16 +591,15 @@ public void testDedupBehaviourOnProjection_SPARK_47217() {
     //Perform leftouter join for exchange table2(firstjoin)
     srcDf = srcDf.join(
             broadcast(table3),
-            srcDf.col("col12").equalTo(
-                table3.col("col32")), "left_outer").
-        select(srcDf.col("col11"), srcDf.col("col12"),
+            srcDf.col("col12").equalTo(table3.col("col32")),
+            "left_outer").select(srcDf.col("col11"),
+            srcDf.col("col12"),
             srcDf.col("col13"),
             table3.col("col33").as("col33_1"));
 
     //Perform left outer joinfor exchangeRateTable1 again(secondjoin)
     Dataset<Row> temp = srcDf.join(broadcast(table3),
-        srcDf.col("col11").equalTo(
-            table3.col("col31")), "left_outer");
+        srcDf.col("col11").equalTo(table3.col("col31")), "left_outer");
 
     srcDf = temp.select(srcDf.col("col11"), srcDf.col("col12"),
         srcDf.col("col13"), srcDf.col("col33_1"),

From 127016c1ad41b9121639b678c9de9080a0ce9d01 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Tue, 5 Mar 2024 21:29:33 -0800
Subject: [PATCH 04/22] SPARK-47217 : Fixing tests and code to try and resolve
 ambiguity in self join conditions

---
 .../scala/org/apache/spark/sql/Dataset.scala  | 44 ++++++++-
 .../apache/spark/sql/JavaDataFrameSuite.java  | 69 -------------
 .../spark/sql/DataFrameAsOfJoinSuite.scala    |  2 +-
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 98 ++++++++++++-------
 4 files changed, 105 insertions(+), 108 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 7a3eae7144b37..d79c2c499bb24 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1150,10 +1150,18 @@ class Dataset[T] private[sql](
 
     // Trigger analysis so in the case of self-join, the analyzer will clone the plan.
     // After the cloning, left and right side will have distinct expression ids.
-    val plan = withPlan(
-      Join(logicalPlan, right.logicalPlan,
-        JoinType(joinType), joinExprs.map(_.expr), JoinHint.NONE))
-      .queryExecution.analyzed.asInstanceOf[Join]
+
+    val plan = try {
+      withPlan(
+        Join(logicalPlan, right.logicalPlan,
+          JoinType(joinType), joinExprs.map(_.expr), JoinHint.NONE))
+        .queryExecution.analyzed.asInstanceOf[Join]
+    } catch {
+      case ae: AnalysisException if ae.message.contains("ambiguous") =>
+        // attempt to resolve ambiguity
+        tryAmbiguityResolution(right, joinExprs, joinType)
+    }
+
 
     // If auto self join alias is disabled, return the plan.
     if (!sparkSession.sessionState.conf.dataFrameSelfJoinAutoResolveAmbiguity) {
@@ -1174,6 +1182,34 @@ class Dataset[T] private[sql](
     JoinWith.resolveSelfJoinCondition(sparkSession.sessionState.analyzer.resolver, plan)
   }
 
+  private def tryAmbiguityResolution(
+      right: Dataset[_],
+      joinExprs: Option[Column],
+      joinType: String) = {
+    val planPart1 = withPlan(
+      Join(logicalPlan, right.logicalPlan,
+        JoinType(joinType), None, JoinHint.NONE))
+      .queryExecution.analyzed.asInstanceOf[Join]
+    val inputSet = planPart1.outputSet
+    val joinExprsRectified = joinExprs.map(_.expr transformUp {
+      case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
+        val attribTagId = attr.metadata.getLong(Dataset.DATASET_ID_KEY)
+        val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
+        val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
+        if (!inputSet.contains(attr) ||
+          (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
+          (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
+          UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
+        } else {
+          attr
+        }
+    })
+    withPlan(
+      Join(planPart1.left, planPart1.right,
+        JoinType(joinType), joinExprsRectified, JoinHint.NONE))
+      .queryExecution.analyzed.asInstanceOf[Join]
+  }
+
   /**
    * Join with another `DataFrame`, using the given join expression. The following performs
    * a full outer join between `df1` and `df2`.
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 09ccbd508e454..94c0a3bd56f1a 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -546,73 +546,4 @@ public void testUDF() {
       .map(row -> row.get(0).toString() + row.getString(1)).toArray(String[]::new);
     Assertions.assertArrayEquals(expected, result);
   }
-
-  @Test
-  public void testDedupBehaviourOnProjection_SPARK_47217() {
-    //Create table1 DF
-    List<Row> table1Data = Arrays.asList(
-      RowFactory.create(1, 2, 3),
-      RowFactory.create(1, 2, 3));
-    StructType table1Scema = new StructType()
-      .add("col11", DataTypes.IntegerType)
-      .add("col12", DataTypes.IntegerType)
-      .add("col13", DataTypes.IntegerType);
-
-    Dataset<Row> table1 = spark.createDataFrame(table1Data, table1Scema);
-
-    //Create table2 DataFrame
-    List<Row> table2Data = Arrays.asList(RowFactory.create(1, 2, 3));
-    StructType table2Schema = new StructType()
-      .add("col21", DataTypes.IntegerType)
-      .add("col22", DataTypes.IntegerType)
-      .add("col23", DataTypes.IntegerType);
-
-    Dataset<Row> table2 = spark.createDataFrame(table2Data, table2Schema);
-
-    //Create table 3 DataFrame
-    List<Row> table3Data = Arrays.asList(RowFactory.create(1, 2, 3), RowFactory.create(1, 2, 3));
-    StructType table3Schema = new StructType()
-      .add("col31", DataTypes.IntegerType)
-      .add("col32", DataTypes.IntegerType)
-      .add("col33", DataTypes.IntegerType);
-
-    Dataset<Row> table3 = spark.createDataFrame(table3Data, table3Schema);
-
-    //Perform left outer join for table2
-    Dataset<Row> srcDf = table1.join(
-        table2,
-        table1.col("col11").equalTo(table2.col("col21")),
-        "left_outer").select(
-        table1.col("col11"),
-        table1.col("col12"),
-        table1.col("col13"),
-        table2.col("col22"));
-
-    //Perform leftouter join for exchange table2(firstjoin)
-    srcDf = srcDf.join(
-            broadcast(table3),
-            srcDf.col("col12").equalTo(table3.col("col32")),
-            "left_outer").select(srcDf.col("col11"),
-            srcDf.col("col12"),
-            srcDf.col("col13"),
-            table3.col("col33").as("col33_1"));
-
-    //Perform left outer joinfor exchangeRateTable1 again(secondjoin)
-    Dataset<Row> temp = srcDf.join(broadcast(table3),
-        srcDf.col("col11").equalTo(table3.col("col31")), "left_outer");
-
-    srcDf = temp.select(srcDf.col("col11"), srcDf.col("col12"),
-        srcDf.col("col13"), srcDf.col("col33_1"),
-        table3.col("col33").as("col33_2"));
-
-    // verify optimized plan creation ok
-    srcDf.queryExecution().optimizedPlan();
-    LogicalPlan lp = srcDf.queryExecution().analyzed();
-    // verify attribute ref resolution is correct, i.e it resolves to the right leg of join
-    AttributeReference refToCheck =
-        (AttributeReference) ((Alias)((Project)lp).projectList().last()).child();
-    AttributeSet compareSet =
-        ((Join) ((Project) lp).child()).right().outputSet();
-    Assertions.assertTrue(compareSet.contains(refToCheck));
-  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
index 3db8d374c4bbc..ec80c782b5b9f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
@@ -175,7 +175,7 @@ class DataFrameAsOfJoinSuite extends QueryTest
     )
   }
 
-  test("SPARK_47217: Dedup of relations can impact projected columns resolution -1") {
+  test("SPARK_47217: Dedup of relations can impact projected columns resolution") {
     val (df1, df2) = prepareForAsOfJoin()
     val join1 = df1.join(df2, df1.col("a") === df2.col("a")).select(df2.col("a"), df1.col("b"),
       df2.col("b"), df1.col("a").as("aa"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index 83d289fcd2073..d14cdc9c42251 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -18,8 +18,8 @@
 package org.apache.spark.sql
 
 import org.apache.spark.api.python.PythonEvalType
-import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, PythonUDF, SortOrder}
-import org.apache.spark.sql.catalyst.plans.logical.{Expand, Generate, ScriptInputOutputSchema, ScriptTransformation, Window => WindowPlan}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, BinaryExpression, PythonUDF, SortOrder}
+import org.apache.spark.sql.catalyst.plans.logical.{Expand, Generate, Join, ScriptInputOutputSchema, ScriptTransformation, Window => WindowPlan}
 import org.apache.spark.sql.expressions.Window
 import org.apache.spark.sql.functions.{count, explode, sum, year}
 import org.apache.spark.sql.internal.SQLConf
@@ -97,7 +97,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assert(e.message.contains("ambiguous"))
   }
 
-  test("SPARK-28344: fail ambiguous self join - column ref in join condition") {
+  test("SPARK-28344: NOT AN ambiguous self join - column ref in join condition") {
     val df1 = spark.range(3)
     val df2 = df1.filter($"id" > 0)
 
@@ -118,29 +118,41 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertAmbiguousSelfJoin(df1.join(df2, df1("id") > df2("id")))
+        val df = df1.join(df2, df1("id") > df2("id"))
+        val join = df.queryExecution.analyzed.asInstanceOf[Join]
+        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
+        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
+        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
     }
   }
 
-  test("SPARK-28344: fail ambiguous self join - Dataset.colRegex as column ref") {
+  test("SPARK-28344: Not AN ambiguous self join - Dataset.colRegex as column ref") {
     val df1 = spark.range(3)
     val df2 = df1.filter($"id" > 0)
 
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertAmbiguousSelfJoin(df1.join(df2, df1.colRegex("id") > df2.colRegex("id")))
+        val df = df1.join(df2, df1.colRegex("id") > df2.colRegex("id"))
+        val join = df.queryExecution.analyzed.asInstanceOf[Join]
+        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
+        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
+        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
     }
   }
 
-  test("SPARK-28344: fail ambiguous self join - Dataset.col with nested field") {
+  test("SPARK-28344: Not An ambiguous self join - Dataset.col with nested field") {
     val df1 = spark.read.json(Seq("""{"a": {"b": 1, "c": 1}}""").toDS())
     val df2 = df1.filter($"a.b" > 0)
 
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertAmbiguousSelfJoin(df1.join(df2, df1("a.b") > df2("a.c")))
+        val df = df1.join(df2, df1("a.b") > df2("a.c"))
+        val join = df.queryExecution.analyzed.asInstanceOf[Join]
+        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
+        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
+        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
     }
   }
 
@@ -293,14 +305,14 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assert(col1DsId !== col2DsId)
   }
 
-  test("SPARK-35454: fail ambiguous self join - toDF") {
+  test("SPARK-35454: Not an ambiguous self join - toDF") {
     val df1 = spark.range(3).toDF()
     val df2 = df1.filter($"id" > 0).toDF()
 
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertAmbiguousSelfJoin(df1.join(df2, df1.col("id") > df2.col("id")))
+        df1.join(df2, df1.col("id") > df2.col("id"))
     }
   }
 
@@ -353,20 +365,20 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     // Test for Project
     val df1 = Seq((1, 2, "A1"), (2, 1, "A2")).toDF("key1", "key2", "value")
     val df2 = df1.filter($"value" === "A2")
-    assertAmbiguousSelfJoin(df1.join(df2, df1("key1") === df2("key2")))
-    assertAmbiguousSelfJoin(df2.join(df1, df1("key1") === df2("key2")))
+    df1.join(df2, df1("key1") === df2("key2"))
+    df2.join(df1, df1("key1") === df2("key2"))
 
     // Test for SerializeFromObject
     val df3 = spark.sparkContext.parallelize(1 to 10).map(x => (x, x)).toDF()
     val df4 = df3.filter($"_1" <=> 0)
-    assertAmbiguousSelfJoin(df3.join(df4, df3("_1") === df4("_2")))
-    assertAmbiguousSelfJoin(df4.join(df3, df3("_1") === df4("_2")))
+    df3.join(df4, df3("_1") === df4("_2"))
+    df4.join(df3, df3("_1") === df4("_2"))
 
     // Test For Aggregate
     val df5 = df1.groupBy($"key1").agg(count($"value") as "count")
     val df6 = df5.filter($"key1" > 0)
-    assertAmbiguousSelfJoin(df5.join(df6, df5("key1") === df6("count")))
-    assertAmbiguousSelfJoin(df6.join(df5, df5("key1") === df6("count")))
+    df5.join(df6, df5("key1") === df6("count"))
+    df6.join(df5, df5("key1") === df6("count"))
 
     // Test for MapInPandas
     val mapInPandasUDF = PythonUDF("mapInPandasUDF", null,
@@ -376,8 +388,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       true)
     val df7 = df1.mapInPandas(mapInPandasUDF)
     val df8 = df7.filter($"x" > 0)
-    assertAmbiguousSelfJoin(df7.join(df8, df7("x") === df8("y")))
-    assertAmbiguousSelfJoin(df8.join(df7, df7("x") === df8("y")))
+    df7.join(df8, df7("x") === df8("y"))
+    df8.join(df7, df7("x") === df8("y"))
 
     // Test for FlatMapGroupsInPandas
     val flatMapGroupsInPandasUDF = PythonUDF("flagMapGroupsInPandasUDF", null,
@@ -387,8 +399,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       true)
     val df9 = df1.groupBy($"key1").flatMapGroupsInPandas(flatMapGroupsInPandasUDF)
     val df10 = df9.filter($"x" > 0)
-    assertAmbiguousSelfJoin(df9.join(df10, df9("x") === df10("y")))
-    assertAmbiguousSelfJoin(df10.join(df9, df9("x") === df10("y")))
+    df9.join(df10, df9("x") === df10("y"))
+    df10.join(df9, df9("x") === df10("y"))
 
     // Test for FlatMapCoGroupsInPandas
     val flatMapCoGroupsInPandasUDF = PythonUDF("flagMapCoGroupsInPandasUDF", null,
@@ -399,22 +411,22 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     val df11 = df1.groupBy($"key1").flatMapCoGroupsInPandas(
       df1.groupBy($"key2"), flatMapCoGroupsInPandasUDF)
     val df12 = df11.filter($"x" > 0)
-    assertAmbiguousSelfJoin(df11.join(df12, df11("x") === df12("y")))
-    assertAmbiguousSelfJoin(df12.join(df11, df11("x") === df12("y")))
+    df11.join(df12, df11("x") === df12("y"))
+    df12.join(df11, df11("x") === df12("y"))
 
     // Test for AttachDistributedSequence
     val df13 = df1.withSequenceColumn("seq")
     val df14 = df13.filter($"value" === "A2")
-    assertAmbiguousSelfJoin(df13.join(df14, df13("key1") === df14("key2")))
-    assertAmbiguousSelfJoin(df14.join(df13, df13("key1") === df14("key2")))
+    df13.join(df14, df13("key1") === df14("key2"))
+    df14.join(df13, df13("key1") === df14("key2"))
 
     // Test for Generate
     // Ensure that the root of the plan is Generate
     val df15 = Seq((1, Seq(1, 2, 3))).toDF("a", "intList").select($"a", explode($"intList"))
       .queryExecution.optimizedPlan.find(_.isInstanceOf[Generate]).get.toDF()
     val df16 = df15.filter($"a" > 0)
-    assertAmbiguousSelfJoin(df15.join(df16, df15("a") === df16("col")))
-    assertAmbiguousSelfJoin(df16.join(df15, df15("a") === df16("col")))
+    df15.join(df16, df15("a") === df16("col"))
+    df16.join(df15, df15("a") === df16("col"))
 
     // Test for Expand
     // Ensure that the root of the plan is Expand
@@ -426,8 +438,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
           AttributeReference("y", IntegerType)()),
         df1.queryExecution.logical).toDF()
     val df18 = df17.filter($"x" > 0)
-    assertAmbiguousSelfJoin(df17.join(df18, df17("x") === df18("y")))
-    assertAmbiguousSelfJoin(df18.join(df17, df17("x") === df18("y")))
+    df17.join(df18, df17("x") === df18("y"))
+    df18.join(df17, df17("x") === df18("y"))
 
     // Test for Window
     val dfWithTS = spark.sql("SELECT timestamp'2021-10-15 01:52:00' time, 1 a, 2 b")
@@ -438,8 +450,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       Seq(SortOrder(dfWithTS("a").expr, Ascending)),
       dfWithTS.queryExecution.logical).toDF()
     val df20 = df19.filter($"a" > 0)
-    assertAmbiguousSelfJoin(df19.join(df20, df19("a") === df20("b")))
-    assertAmbiguousSelfJoin(df20.join(df19, df19("a") === df20("b")))
+    df19.join(df20, df19("a") === df20("b"))
+    df20.join(df19, df19("a") === df20("b"))
 
     // Test for ScriptTransformation
     val ioSchema =
@@ -464,8 +476,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       df1.queryExecution.logical,
       ioSchema).toDF()
     val df22 = df21.filter($"x" > 0)
-    assertAmbiguousSelfJoin(df21.join(df22, df21("x") === df22("y")))
-    assertAmbiguousSelfJoin(df22.join(df21, df21("x") === df22("y")))
+    df21.join(df22, df21("x") === df22("y"))
+    df22.join(df21, df21("x") === df22("y"))
   }
 
   test("SPARK-35937: GetDateFieldOperations should skip unresolved nodes") {
@@ -489,13 +501,31 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     }
   }
 
-  test("SPARK_47217: Dedup of relations can impact projected columns resolution") {
+  test("SPARK_47217: deduplication of project causes ambiguity in resolution") {
     val df = Seq((1, 2)).toDF("a", "b")
     val df2 = df.select(df("a").as("aa"), df("b").as("bb"))
     val df3 = df2.join(df, df2("bb") === df("b")).select(df2("aa"), df("a"))
-
     checkAnswer(
       df3,
       Row(1, 1) :: Nil)
   }
+
+  test("SPARK-47217. deduplication in nested joins focussing on projection") {
+    val df1 = Seq((1, 2)).toDF("a", "b")
+    val df2 = Seq((1, 2)).toDF("aa", "bb")
+    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
+      df2("aa"), df1("b"))
+    val df3 = df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).
+      select(df1Joindf2("aa"), df1("a"))
+    df3.queryExecution.assertAnalyzed()
+  }
+
+  test("SPARK-47217. deduplication in nested joins focusing on condition") {
+    val df1 = Seq((1, 2)).toDF("a", "b")
+    val df2 = Seq((1, 2)).toDF("aa", "bb")
+    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"),
+      df2("aa"), df1("b"))
+    val df3 = df1Joindf2.join(df1, df1Joindf2("aa") === df1("a"))
+    df3.queryExecution.assertAnalyzed()
+  }
 }

From b8e369c239f17cba34df2c816a6f5b412c5aa63b Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Tue, 5 Mar 2024 23:50:18 -0800
Subject: [PATCH 05/22] SPARK-47217 : Fix unused import issue

---
 .../java/test/org/apache/spark/sql/JavaDataFrameSuite.java  | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 94c0a3bd56f1a..26a19cbed1b91 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -24,7 +24,6 @@
 import java.math.BigInteger;
 import java.math.BigDecimal;
 
-import org.apache.spark.sql.catalyst.plans.logical.Join;
 import scala.collection.Seq;
 import scala.jdk.javaapi.CollectionConverters;
 
@@ -32,11 +31,6 @@
 import com.google.common.primitives.Ints;
 import org.junit.jupiter.api.*;
 
-import org.apache.spark.sql.catalyst.expressions.Alias;
-import org.apache.spark.sql.catalyst.expressions.AttributeReference;
-import org.apache.spark.sql.catalyst.expressions.AttributeSet;
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
-import org.apache.spark.sql.catalyst.plans.logical.Project;
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.Column;

From 872fece15d69b27a1bd40c7ea1e99ce2a604788c Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Wed, 6 Mar 2024 14:36:19 -0800
Subject: [PATCH 06/22] SPARK-47217 : fixed bug and made assertions in existing
 tests for correct resolution of attributes

---
 .../analysis/ColumnResolutionHelper.scala     |  19 +--
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 148 +++++++++++-------
 2 files changed, 99 insertions(+), 68 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index 22361b0e73c77..bd1e561dfe062 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -485,26 +485,15 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
 
       resolveOnDatasetId = (datasetid: Long, name: String) => {
         def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[LogicalPlan] = {
-          var currentLp = lp
-          while(currentLp.children.size < 2) {
-            if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
-              return Option(currentLp)
-            } else {
-              if (currentLp.children.size == 1) {
-                currentLp = currentLp.children.head
-              } else {
-                // leaf node
-                return None
-              }
-            }
+          if (lp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
+             Option(lp)
+          } else {
+            None
           }
-          None
         }
-
         val binaryNodeOpt = q.collectFirst {
           case bn: BinaryNode => bn
         }
-
         val resolveOnAttribs = binaryNodeOpt match {
           case Some(bn) =>
             val leftDefOpt = findUnaryNodeMatchingTagId(bn.left)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index d14cdc9c42251..c962f31eefacb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -97,6 +97,27 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assert(e.message.contains("ambiguous"))
   }
 
+  private def assertCorrectResolution(
+      df: => DataFrame,
+      leftResolution: Resolution.Resolution,
+      rightResolution: Resolution.Resolution): Unit = {
+    val join = df.queryExecution.analyzed.asInstanceOf[Join]
+    val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
+    leftResolution match {
+      case Resolution.LeftConditionToLeftLeg =>
+        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
+      case Resolution.LeftConditionToRightLeg =>
+        assert(join.right.outputSet.contains(binaryCondition.left.references.head))
+    }
+
+    rightResolution match {
+      case Resolution.RightConditionToLeftLeg =>
+        assert(join.left.outputSet.contains(binaryCondition.right.references.head))
+      case Resolution.RightConditionToRightLeg =>
+        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
+    }
+  }
+
   test("SPARK-28344: NOT AN ambiguous self join - column ref in join condition") {
     val df1 = spark.range(3)
     val df2 = df1.filter($"id" > 0)
@@ -118,11 +139,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-        val df = df1.join(df2, df1("id") > df2("id"))
-        val join = df.queryExecution.analyzed.asInstanceOf[Join]
-        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
-        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
-        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
+        assertCorrectResolution(df1.join(df2, df1("id") > df2("id")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     }
   }
 
@@ -133,11 +151,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-        val df = df1.join(df2, df1.colRegex("id") > df2.colRegex("id"))
-        val join = df.queryExecution.analyzed.asInstanceOf[Join]
-        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
-        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
-        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
+      assertCorrectResolution(df1.join(df2, df1.colRegex("id") > df2.colRegex("id")),
+        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     }
   }
 
@@ -148,11 +163,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-        val df = df1.join(df2, df1("a.b") > df2("a.c"))
-        val join = df.queryExecution.analyzed.asInstanceOf[Join]
-        val binaryCondition = join.condition.get.asInstanceOf[BinaryExpression]
-        assert(join.left.outputSet.contains(binaryCondition.left.references.head))
-        assert(join.right.outputSet.contains(binaryCondition.right.references.head))
+      assertCorrectResolution( df1.join(df2, df1("a.b") > df2("a.c")),
+        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     }
   }
 
@@ -312,7 +324,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-        df1.join(df2, df1.col("id") > df2.col("id"))
+      assertCorrectResolution(df1.join(df2, df1.col("id") > df2.col("id")),
+        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     }
   }
 
@@ -363,22 +376,30 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   test("SPARK-36874: DeduplicateRelations should copy dataset_id tag " +
     "to avoid ambiguous self join") {
     // Test for Project
+
     val df1 = Seq((1, 2, "A1"), (2, 1, "A2")).toDF("key1", "key2", "value")
     val df2 = df1.filter($"value" === "A2")
-    df1.join(df2, df1("key1") === df2("key2"))
-    df2.join(df1, df1("key1") === df2("key2"))
+ /*   assertCorrectResolution(df1.join(df2, df1("key1") === df2("key2")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df2.join(df1, df1("key1") === df2("key2")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
 
     // Test for SerializeFromObject
     val df3 = spark.sparkContext.parallelize(1 to 10).map(x => (x, x)).toDF()
     val df4 = df3.filter($"_1" <=> 0)
-    df3.join(df4, df3("_1") === df4("_2"))
-    df4.join(df3, df3("_1") === df4("_2"))
+    assertCorrectResolution(df3.join(df4, df3("_1") === df4("_2")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df4.join(df3, df3("_1") === df4("_2")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
     // Test For Aggregate
     val df5 = df1.groupBy($"key1").agg(count($"value") as "count")
     val df6 = df5.filter($"key1" > 0)
-    df5.join(df6, df5("key1") === df6("count"))
-    df6.join(df5, df5("key1") === df6("count"))
+    assertCorrectResolution(df5.join(df6, df5("key1") === df6("count")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df6.join(df5, df5("key1") === df6("count")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
     // Test for MapInPandas
     val mapInPandasUDF = PythonUDF("mapInPandasUDF", null,
@@ -388,8 +409,10 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       true)
     val df7 = df1.mapInPandas(mapInPandasUDF)
     val df8 = df7.filter($"x" > 0)
-    df7.join(df8, df7("x") === df8("y"))
-    df8.join(df7, df7("x") === df8("y"))
+    assertCorrectResolution(df7.join(df8, df7("x") === df8("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df8.join(df7, df7("x") === df8("y")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
     // Test for FlatMapGroupsInPandas
     val flatMapGroupsInPandasUDF = PythonUDF("flagMapGroupsInPandasUDF", null,
@@ -399,9 +422,11 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       true)
     val df9 = df1.groupBy($"key1").flatMapGroupsInPandas(flatMapGroupsInPandasUDF)
     val df10 = df9.filter($"x" > 0)
-    df9.join(df10, df9("x") === df10("y"))
-    df10.join(df9, df9("x") === df10("y"))
-
+    assertCorrectResolution(df9.join(df10, df9("x") === df10("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df10.join(df9, df9("x") === df10("y")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+ */
     // Test for FlatMapCoGroupsInPandas
     val flatMapCoGroupsInPandasUDF = PythonUDF("flagMapCoGroupsInPandasUDF", null,
       StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
@@ -411,22 +436,27 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     val df11 = df1.groupBy($"key1").flatMapCoGroupsInPandas(
       df1.groupBy($"key2"), flatMapCoGroupsInPandasUDF)
     val df12 = df11.filter($"x" > 0)
-    df11.join(df12, df11("x") === df12("y"))
-    df12.join(df11, df11("x") === df12("y"))
+  /*  assertCorrectResolution(df11.join(df12, df11("x") === df12("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg) */
+    assertCorrectResolution(df12.join(df11, df11("x") === df12("y")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
     // Test for AttachDistributedSequence
     val df13 = df1.withSequenceColumn("seq")
     val df14 = df13.filter($"value" === "A2")
-    df13.join(df14, df13("key1") === df14("key2"))
-    df14.join(df13, df13("key1") === df14("key2"))
-
+    assertCorrectResolution(df13.join(df14, df13("key1") === df14("key2")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df14.join(df13, df13("key1") === df14("key2")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
     // Test for Generate
     // Ensure that the root of the plan is Generate
     val df15 = Seq((1, Seq(1, 2, 3))).toDF("a", "intList").select($"a", explode($"intList"))
       .queryExecution.optimizedPlan.find(_.isInstanceOf[Generate]).get.toDF()
     val df16 = df15.filter($"a" > 0)
-    df15.join(df16, df15("a") === df16("col"))
-    df16.join(df15, df15("a") === df16("col"))
+    assertCorrectResolution(df15.join(df16, df15("a") === df16("col")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df16.join(df15, df15("a") === df16("col")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
     // Test for Expand
     // Ensure that the root of the plan is Expand
@@ -438,9 +468,10 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
           AttributeReference("y", IntegerType)()),
         df1.queryExecution.logical).toDF()
     val df18 = df17.filter($"x" > 0)
-    df17.join(df18, df17("x") === df18("y"))
-    df18.join(df17, df17("x") === df18("y"))
-
+    assertCorrectResolution(df17.join(df18, df17("x") === df18("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df18.join(df17, df17("x") === df18("y")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
     // Test for Window
     val dfWithTS = spark.sql("SELECT timestamp'2021-10-15 01:52:00' time, 1 a, 2 b")
     // Ensure that the root of the plan is Window
@@ -450,9 +481,10 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       Seq(SortOrder(dfWithTS("a").expr, Ascending)),
       dfWithTS.queryExecution.logical).toDF()
     val df20 = df19.filter($"a" > 0)
-    df19.join(df20, df19("a") === df20("b"))
-    df20.join(df19, df19("a") === df20("b"))
-
+    assertCorrectResolution(df19.join(df20, df19("a") === df20("b")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df20.join(df19, df19("a") === df20("b")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
     // Test for ScriptTransformation
     val ioSchema =
       ScriptInputOutputSchema(
@@ -476,8 +508,10 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       df1.queryExecution.logical,
       ioSchema).toDF()
     val df22 = df21.filter($"x" > 0)
-    df21.join(df22, df21("x") === df22("y"))
-    df22.join(df21, df21("x") === df22("y"))
+    assertCorrectResolution(df21.join(df22, df21("x") === df22("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+    assertCorrectResolution(df22.join(df21, df21("x") === df22("y")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
   }
 
   test("SPARK-35937: GetDateFieldOperations should skip unresolved nodes") {
@@ -515,17 +549,25 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     val df2 = Seq((1, 2)).toDF("aa", "bb")
     val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
       df2("aa"), df1("b"))
-    val df3 = df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).
-      select(df1Joindf2("aa"), df1("a"))
-    df3.queryExecution.assertAnalyzed()
-  }
 
-  test("SPARK-47217. deduplication in nested joins focusing on condition") {
-    val df1 = Seq((1, 2)).toDF("a", "b")
-    val df2 = Seq((1, 2)).toDF("aa", "bb")
-    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"),
-      df2("aa"), df1("b"))
-    val df3 = df1Joindf2.join(df1, df1Joindf2("aa") === df1("a"))
-    df3.queryExecution.assertAnalyzed()
+    assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+
+    assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+    df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"), df1("a")).
+      queryExecution.analyzed
+
+    df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"), df1("a")).
+      queryExecution.analyzed
   }
 }
+
+object Resolution extends Enumeration {
+  type Resolution = Value
+
+  val LeftConditionToLeftLeg, LeftConditionToRightLeg, RightConditionToRightLeg,
+  RightConditionToLeftLeg = Value
+}
+

From 8ed6aa48502dd4ca820671b58a7136cb13fbdcb7 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Wed, 6 Mar 2024 17:30:13 -0800
Subject: [PATCH 07/22] SPARK-47217 : added more assetions

---
 .../spark/sql/DataFrameSelfJoinSuite.scala     | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index e0388f12cee67..b6bf214e84075 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql
 
 import org.apache.spark.api.python.PythonEvalType
 import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, BinaryExpression, PythonUDF, SortOrder}
-import org.apache.spark.sql.catalyst.plans.logical.{Expand, Generate, Join, ScriptInputOutputSchema, ScriptTransformation, Window => WindowPlan}
+import org.apache.spark.sql.catalyst.plans.logical.{Expand, Generate, Join, Project, ScriptInputOutputSchema, ScriptTransformation, Window => WindowPlan}
 import org.apache.spark.sql.expressions.Window
 import org.apache.spark.sql.functions.{count, explode, sum, year}
 import org.apache.spark.sql.internal.SQLConf
@@ -566,11 +566,17 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")),
       Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
-    df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"), df1("a")).
-      queryExecution.analyzed
-
-    df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"), df1("a")).
-      queryExecution.analyzed
+    val proj1 = df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
+      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+    val join1 = proj1.child.asInstanceOf[Join]
+    assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
+    assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+
+    val proj2 = df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
+      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+    val join2 = proj2.child.asInstanceOf[Join]
+    assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
+    assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
   }
 }
 

From 6b3b1d4ed0549a627df837631118d6058ec5f91a Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Wed, 6 Mar 2024 21:24:26 -0800
Subject: [PATCH 08/22] SPARK-47217 : fixed a bug and uncommented tests which
 were inadvertently commented

---
 .../analysis/ColumnResolutionHelper.scala     | 39 +++++++++++++++----
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 10 ++---
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index bd1e561dfe062..4aac2c6c70676 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -484,23 +484,48 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       },
 
       resolveOnDatasetId = (datasetid: Long, name: String) => {
-        def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[LogicalPlan] = {
-          if (lp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
-             Option(lp)
-          } else {
-            None
+        def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[(LogicalPlan, Int)] = {
+          var currentLp = lp
+          var depth = 0
+          while(true) {
+            if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
+              return Option(currentLp, depth)
+            } else {
+              if (currentLp.children.size == 1) {
+                currentLp = currentLp.children.head
+              } else {
+                // leaf node or node is a binary node
+                return None
+              }
+            }
+            depth += 1
           }
+          None
         }
+
         val binaryNodeOpt = q.collectFirst {
           case bn: BinaryNode => bn
         }
+
         val resolveOnAttribs = binaryNodeOpt match {
           case Some(bn) =>
             val leftDefOpt = findUnaryNodeMatchingTagId(bn.left)
             val rightDefOpt = findUnaryNodeMatchingTagId(bn.right)
             (leftDefOpt, rightDefOpt) match {
-              case (None, Some(lp)) => lp.output
-              case (Some(lp), None) => lp.output
+
+              case (None, Some((lp, _))) => lp.output
+
+              case (Some((lp, _)), None) => lp.output
+
+              case  (Some((lp1, depth1)), Some((lp2, depth2))) =>
+                if (depth1 == depth2) {
+                  q.children.head.output
+                } else if (depth1 < depth2) {
+                  lp1.output
+                } else {
+                  lp2.output
+                }
+
               case _ => q.children.head.output
             }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index b6bf214e84075..9aa7707f83617 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -379,7 +379,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
 
     val df1 = Seq((1, 2, "A1"), (2, 1, "A2")).toDF("key1", "key2", "value")
     val df2 = df1.filter($"value" === "A2")
- /*   assertCorrectResolution(df1.join(df2, df1("key1") === df2("key2")),
+    assertCorrectResolution(df1.join(df2, df1("key1") === df2("key2")),
       Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     assertCorrectResolution(df2.join(df1, df1("key1") === df2("key2")),
       Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
@@ -426,7 +426,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     assertCorrectResolution(df10.join(df9, df9("x") === df10("y")),
       Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
- */
+
     // Test for FlatMapCoGroupsInPandas
     val flatMapCoGroupsInPandasUDF = PythonUDF("flagMapCoGroupsInPandasUDF", null,
       StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
@@ -436,8 +436,8 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     val df11 = df1.groupBy($"key1").flatMapCoGroupsInPandas(
       df1.groupBy($"key2"), flatMapCoGroupsInPandasUDF)
     val df12 = df11.filter($"x" > 0)
-  /*  assertCorrectResolution(df11.join(df12, df11("x") === df12("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg) */
+    assertCorrectResolution(df11.join(df12, df11("x") === df12("y")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
     assertCorrectResolution(df12.join(df11, df11("x") === df12("y")),
       Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
@@ -584,6 +584,6 @@ object Resolution extends Enumeration {
   type Resolution = Value
 
   val LeftConditionToLeftLeg, LeftConditionToRightLeg, RightConditionToRightLeg,
-  RightConditionToLeftLeg = Value
+    RightConditionToLeftLeg = Value
 }
 

From f9653ec7d439776a8b3c384fe1bebf045e7ad9ca Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 7 Mar 2024 15:20:25 -0800
Subject: [PATCH 09/22] SPARK-47217 : added more tests and fixed inconsistency

---
 .../scala/org/apache/spark/sql/Dataset.scala  | 54 +++++++++++--------
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 49 +++++++++++++++--
 2 files changed, 77 insertions(+), 26 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index d79c2c499bb24..2bd48e45dd846 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -38,6 +38,7 @@ import org.apache.spark.api.r.RRDD
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.resource.ResourceProfile
+import org.apache.spark.sql.Dataset.DATASET_ID_KEY
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, QueryPlanningTracker, ScalaReflection, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
@@ -1151,17 +1152,9 @@ class Dataset[T] private[sql](
     // Trigger analysis so in the case of self-join, the analyzer will clone the plan.
     // After the cloning, left and right side will have distinct expression ids.
 
-    val plan = try {
-      withPlan(
-        Join(logicalPlan, right.logicalPlan,
-          JoinType(joinType), joinExprs.map(_.expr), JoinHint.NONE))
-        .queryExecution.analyzed.asInstanceOf[Join]
-    } catch {
-      case ae: AnalysisException if ae.message.contains("ambiguous") =>
-        // attempt to resolve ambiguity
-        tryAmbiguityResolution(right, joinExprs, joinType)
-    }
-
+    val plan = withPlan(
+      tryAmbiguityResolution(right, joinExprs, joinType)
+    ).queryExecution.analyzed.asInstanceOf[Join]
 
     // If auto self join alias is disabled, return the plan.
     if (!sparkSession.sessionState.conf.dataFrameSelfJoinAutoResolveAmbiguity) {
@@ -1199,15 +1192,13 @@ class Dataset[T] private[sql](
         if (!inputSet.contains(attr) ||
           (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
           (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
-          UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
+          UnresolvedAttributeWithTag(attr, attribTagId)
         } else {
           attr
         }
     })
-    withPlan(
-      Join(planPart1.left, planPart1.right,
-        JoinType(joinType), joinExprsRectified, JoinHint.NONE))
-      .queryExecution.analyzed.asInstanceOf[Join]
+
+    Join(planPart1.left, planPart1.right, JoinType(joinType), joinExprsRectified, JoinHint.NONE)
   }
 
   /**
@@ -1624,11 +1615,25 @@ class Dataset[T] private[sql](
     val inputSet = logicalPlan.outputSet
     val rectifiedNamedExprs = namedExprs.map(ne => ne match {
 
-      case al: Alias if !al.references.subsetOf(inputSet) &&
+      case al: Alias if (!al.references.subsetOf(inputSet) || al.references.exists(attr =>
+        attr.metadata.contains(DATASET_ID_KEY) && attr.metadata.getLong(DATASET_ID_KEY) !=
+          inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
+            if (x.metadata.contains(DATASET_ID_KEY)) {
+              x.metadata.getLong(DATASET_ID_KEY)
+            } else {
+              -1
+          }).get)) &&
         al.nonInheritableMetadataKeys.contains(Dataset.DATASET_ID_KEY) =>
         val unresolvedExpr = al.child.transformUp {
-          case attr: AttributeReference if !inputSet.contains(attr) &&
-            attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
+          case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) &&
+            (!inputSet.contains(attr) || attr.metadata.getLong(DATASET_ID_KEY) !=
+              inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
+                if (x.metadata.contains(DATASET_ID_KEY)) {
+                  x.metadata.getLong(DATASET_ID_KEY)
+                } else {
+                  -1
+                }).get)
+             =>
             UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
         }
         val newAl = al.copy(child = unresolvedExpr, name = al.name)(exprId = al.exprId,
@@ -1637,8 +1642,15 @@ class Dataset[T] private[sql](
         newAl.copyTagsFrom(al)
         newAl
 
-      case attr: Attribute if !inputSet.contains(attr) &&
-        attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
+      case attr: Attribute if attr.metadata.contains(Dataset.DATASET_ID_KEY) &&
+        (!inputSet.contains(attr) || attr.metadata.getLong(DATASET_ID_KEY) !=
+          inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
+            if (x.metadata.contains(DATASET_ID_KEY)) {
+              x.metadata.getLong(DATASET_ID_KEY)
+            } else {
+              -1
+            }).get)
+         =>
         UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
 
       case _ => ne
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index 9aa7707f83617..4fccf9d2415cb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -189,7 +189,9 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertAmbiguousSelfJoin(df1.join(df2).select(df2("id")))
+      val proj1 = df1.join(df2).select(df2("id")).queryExecution.analyzed.asInstanceOf[Project]
+      val join1 = proj1.child.asInstanceOf[Join]
+      assert(proj1.projectList(0).references.subsetOf(join1.right.outputSet))
     }
   }
 
@@ -229,7 +231,11 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
       assertAmbiguousSelfJoin(df1.join(df2).join(df3, df2("id") < df3("id")))
-      assertAmbiguousSelfJoin(df1.join(df4).join(df2).select(df2("id")))
+
+      val proj1 = df1.join(df4).join(df2).select(df2("id")).queryExecution.analyzed.
+        asInstanceOf[Project]
+      val join1 = proj1.child.asInstanceOf[Join]
+      assert(proj1.projectList(0).references.subsetOf(join1.right.outputSet))
     }
   }
 
@@ -261,8 +267,17 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       TestData(2, "personnel"),
       TestData(3, "develop")).toDS()
     val emp3 = emp1.join(emp2, emp1("key") === emp2("key")).select(emp1("*"))
-    assertAmbiguousSelfJoin(emp1.join(emp3, emp1.col("key") === emp3.col("key"),
-      "left_outer").select(emp1.col("*"), emp3.col("key").as("e2")))
+
+    assertCorrectResolution(emp1.join(emp3, emp1.col("key") === emp3.col("key")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+
+    val proj1 = emp1.join(emp3, emp1.col("key") === emp3.col("key"),
+      "left_outer").select(emp1.col("*"), emp3.col("key").as("e2")).
+      queryExecution.analyzed.asInstanceOf[Project]
+    val join1 = proj1.child.asInstanceOf[Join]
+    assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
+    assert(proj1.projectList(1).references.subsetOf(join1.left.outputSet))
+    assert(proj1.projectList(2).references.subsetOf(join1.right.outputSet))
   }
 
   test("df.show() should also not change dataset_id of LogicalPlan") {
@@ -554,7 +569,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       Row(1, 1) :: Nil)
   }
 
-  test("SPARK-47217. deduplication in nested joins") {
+  test("SPARK-47217. deduplication in nested joins with join attribute aliased") {
     val df1 = Seq((1, 2)).toDF("a", "b")
     val df2 = Seq((1, 2)).toDF("aa", "bb")
     val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
@@ -578,6 +593,30 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
     assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
   }
+
+  test("SPARK-47217. deduplication in nested joins without join attribute aliased") {
+    val df1 = Seq((1, 2)).toDF("a", "b")
+    val df2 = Seq((1, 2)).toDF("aa", "bb")
+    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"), df2("aa"), df1("b"))
+
+    assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("a") === df1("a")),
+      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+
+    assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("a") === df1("a")),
+      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+    val proj1 = df1Joindf2.join(df1, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
+      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+    val join1 = proj1.child.asInstanceOf[Join]
+    assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
+    assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+
+    val proj2 = df1.join(df1Joindf2, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
+      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+    val join2 = proj2.child.asInstanceOf[Join]
+    assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
+    assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
+  }
 }
 
 object Resolution extends Enumeration {

From 7150c9887acd217bdae3c9e58d58d37e070ab0e1 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 8 Mar 2024 11:49:22 -0800
Subject: [PATCH 10/22] SPARK-47217 : fixed test failure

---
 .../main/scala/org/apache/spark/sql/Dataset.scala    | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 2bd48e45dd846..0f15fcf51b8f3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1615,15 +1615,15 @@ class Dataset[T] private[sql](
     val inputSet = logicalPlan.outputSet
     val rectifiedNamedExprs = namedExprs.map(ne => ne match {
 
-      case al: Alias if (!al.references.subsetOf(inputSet) || al.references.exists(attr =>
+      case al: Alias if !al.references.subsetOf(inputSet) || al.references.exists(attr =>
         attr.metadata.contains(DATASET_ID_KEY) && attr.metadata.getLong(DATASET_ID_KEY) !=
           inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
             if (x.metadata.contains(DATASET_ID_KEY)) {
               x.metadata.getLong(DATASET_ID_KEY)
             } else {
-              -1
-          }).get)) &&
-        al.nonInheritableMetadataKeys.contains(Dataset.DATASET_ID_KEY) =>
+              Dataset.this.id
+          }).get)
+         =>
         val unresolvedExpr = al.child.transformUp {
           case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) &&
             (!inputSet.contains(attr) || attr.metadata.getLong(DATASET_ID_KEY) !=
@@ -1631,7 +1631,7 @@ class Dataset[T] private[sql](
                 if (x.metadata.contains(DATASET_ID_KEY)) {
                   x.metadata.getLong(DATASET_ID_KEY)
                 } else {
-                  -1
+                  Dataset.this.id
                 }).get)
              =>
             UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
@@ -1648,7 +1648,7 @@ class Dataset[T] private[sql](
             if (x.metadata.contains(DATASET_ID_KEY)) {
               x.metadata.getLong(DATASET_ID_KEY)
             } else {
-              -1
+              Dataset.this.id
             }).get)
          =>
         UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))

From 6e03ae0aa625061648b49284756fcb6bc6111cf0 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Mon, 11 Mar 2024 13:58:07 -0700
Subject: [PATCH 11/22] SPARK-47320 : incorporate review comments

---
 .../scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala | 2 +-
 .../scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
index ec80c782b5b9f..4d1953f21890f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAsOfJoinSuite.scala
@@ -175,7 +175,7 @@ class DataFrameAsOfJoinSuite extends QueryTest
     )
   }
 
-  test("SPARK_47217: Dedup of relations can impact projected columns resolution") {
+  test("SPARK-47217: Dedup of relations can impact projected columns resolution") {
     val (df1, df2) = prepareForAsOfJoin()
     val join1 = df1.join(df2, df1.col("a") === df2.col("a")).select(df2.col("a"), df1.col("b"),
       df2.col("b"), df1.col("a").as("aa"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index 4fccf9d2415cb..32a6deb436f9f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -560,7 +560,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     }
   }
 
-  test("SPARK_47217: deduplication of project causes ambiguity in resolution") {
+  test("SPARK-47217: deduplication of project causes ambiguity in resolution") {
     val df = Seq((1, 2)).toDF("a", "b")
     val df2 = df.select(df("a").as("aa"), df("b").as("bb"))
     val df3 = df2.join(df, df2("bb") === df("b")).select(df2("aa"), df("a"))
@@ -569,7 +569,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
       Row(1, 1) :: Nil)
   }
 
-  test("SPARK-47217. deduplication in nested joins with join attribute aliased") {
+  test("SPARK-47217: deduplication in nested joins with join attribute aliased") {
     val df1 = Seq((1, 2)).toDF("a", "b")
     val df2 = Seq((1, 2)).toDF("aa", "bb")
     val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
@@ -594,7 +594,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
   }
 
-  test("SPARK-47217. deduplication in nested joins without join attribute aliased") {
+  test("SPARK-47217: deduplication in nested joins without join attribute aliased") {
     val df1 = Seq((1, 2)).toDF("a", "b")
     val df2 = Seq((1, 2)).toDF("aa", "bb")
     val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"), df2("aa"), df1("b"))

From d2175d49afd5f997725aaf97e197fa34f4e77634 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 14 Mar 2024 12:20:49 -0700
Subject: [PATCH 12/22] SPARK-47320: reverting earlier change and refactoring

---
 .../scala/org/apache/spark/sql/Dataset.scala  | 83 ++++++-------------
 1 file changed, 27 insertions(+), 56 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 0f15fcf51b8f3..82e3f01b7189c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1185,19 +1185,8 @@ class Dataset[T] private[sql](
       .queryExecution.analyzed.asInstanceOf[Join]
     val inputSet = planPart1.outputSet
     val joinExprsRectified = joinExprs.map(_.expr transformUp {
-      case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        val attribTagId = attr.metadata.getLong(Dataset.DATASET_ID_KEY)
-        val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
-        val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
-        if (!inputSet.contains(attr) ||
-          (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
-          (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
-          UnresolvedAttributeWithTag(attr, attribTagId)
-        } else {
-          attr
-        }
+      case attr: AttributeReference => convertToUnresolvedIfNeeded(attr, inputSet)
     })
-
     Join(planPart1.left, planPart1.right, JoinType(joinType), joinExprsRectified, JoinHint.NONE)
   }
 
@@ -1208,7 +1197,7 @@ class Dataset[T] private[sql](
    * {{{
    *   // Scala:
    *   import org.apache.spark.sql.functions._
-   *   df1.join(df2, $"df1Key" === $"df2Key", "outer")
+   *   df1.join(df2, $"df1Key" === $"df2Key", "outer"
    *
    *   // Java:
    *   import static org.apache.spark.sql.functions.*;
@@ -1613,49 +1602,9 @@ class Dataset[T] private[sql](
     }
     val namedExprs = untypedCols.map(_.named)
     val inputSet = logicalPlan.outputSet
-    val rectifiedNamedExprs = namedExprs.map(ne => ne match {
-
-      case al: Alias if !al.references.subsetOf(inputSet) || al.references.exists(attr =>
-        attr.metadata.contains(DATASET_ID_KEY) && attr.metadata.getLong(DATASET_ID_KEY) !=
-          inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
-            if (x.metadata.contains(DATASET_ID_KEY)) {
-              x.metadata.getLong(DATASET_ID_KEY)
-            } else {
-              Dataset.this.id
-          }).get)
-         =>
-        val unresolvedExpr = al.child.transformUp {
-          case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) &&
-            (!inputSet.contains(attr) || attr.metadata.getLong(DATASET_ID_KEY) !=
-              inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
-                if (x.metadata.contains(DATASET_ID_KEY)) {
-                  x.metadata.getLong(DATASET_ID_KEY)
-                } else {
-                  Dataset.this.id
-                }).get)
-             =>
-            UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
-        }
-        val newAl = al.copy(child = unresolvedExpr, name = al.name)(exprId = al.exprId,
-          qualifier = al.qualifier, explicitMetadata = al.explicitMetadata,
-          nonInheritableMetadataKeys = al.nonInheritableMetadataKeys)
-        newAl.copyTagsFrom(al)
-        newAl
-
-      case attr: Attribute if attr.metadata.contains(Dataset.DATASET_ID_KEY) &&
-        (!inputSet.contains(attr) || attr.metadata.getLong(DATASET_ID_KEY) !=
-          inputSet.find(_.canonicalized == attr.canonicalized).map(x =>
-            if (x.metadata.contains(DATASET_ID_KEY)) {
-              x.metadata.getLong(DATASET_ID_KEY)
-            } else {
-              Dataset.this.id
-            }).get)
-         =>
-        UnresolvedAttributeWithTag(attr, attr.metadata.getLong(Dataset.DATASET_ID_KEY))
-
-      case _ => ne
-
-    })
+    val rectifiedNamedExprs = namedExprs.map(ne => (ne transformUp {
+      case attr: AttributeReference => convertToUnresolvedIfNeeded(attr, inputSet)
+    }).asInstanceOf[NamedExpression])
     Project(rectifiedNamedExprs, logicalPlan)
   }
 
@@ -4316,6 +4265,28 @@ class Dataset[T] private[sql](
     queryExecution.analyzed.semanticHash()
   }
 
+  private def convertToUnresolvedIfNeeded(attr: AttributeReference, inputSet: AttributeSet):
+  NamedExpression = {
+    val attrDatasetIdOpt = if (attr.metadata.contains(Dataset.DATASET_ID_KEY)) {
+      Option(attr.metadata.getLong(DATASET_ID_KEY))
+    } else {
+      None
+    }
+    if (attrDatasetIdOpt.isDefined && inputSet.find(_ == attr).forall(x => {
+      val datasetIdX = if (x.metadata.contains(DATASET_ID_KEY)) {
+        x.metadata.getLong(DATASET_ID_KEY)
+      } else {
+        Dataset.this.id
+      }
+      attrDatasetIdOpt.get != datasetIdX
+    })) {
+      UnresolvedAttributeWithTag(attr, attrDatasetIdOpt.get)
+    } else {
+      attr
+    }
+  }
+
+
   ////////////////////////////////////////////////////////////////////////////
   // For Python API
   ////////////////////////////////////////////////////////////////////////////

From f3d280e2bf3ec1539ffa472d28e0f27760e87c18 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 14 Mar 2024 13:23:12 -0700
Subject: [PATCH 13/22] SPARK-47320: refactoring

---
 .../src/main/scala/org/apache/spark/sql/Dataset.scala | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 82e3f01b7189c..1feff155b8706 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1151,7 +1151,6 @@ class Dataset[T] private[sql](
 
     // Trigger analysis so in the case of self-join, the analyzer will clone the plan.
     // After the cloning, left and right side will have distinct expression ids.
-
     val plan = withPlan(
       tryAmbiguityResolution(right, joinExprs, joinType)
     ).queryExecution.analyzed.asInstanceOf[Join]
@@ -1328,7 +1327,6 @@ class Dataset[T] private[sql](
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
         UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
     }
-
     val rightAsOfExpr = rightAsOf.expr.transformUp {
       case a: AttributeReference if other.logicalPlan.outputSet.contains(a) =>
         val index = other.logicalPlan.output.indexWhere(_.exprId == a.exprId)
@@ -1337,7 +1335,6 @@ class Dataset[T] private[sql](
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
         UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
     }
-
     withPlan {
       AsOfJoin(
         joined.left, joined.right,
@@ -1600,12 +1597,11 @@ class Dataset[T] private[sql](
 
       case other => other
     }
-    val namedExprs = untypedCols.map(_.named)
     val inputSet = logicalPlan.outputSet
-    val rectifiedNamedExprs = namedExprs.map(ne => (ne transformUp {
+    val namedExprs = untypedCols.map(ne => (ne.named transformUp {
       case attr: AttributeReference => convertToUnresolvedIfNeeded(attr, inputSet)
     }).asInstanceOf[NamedExpression])
-    Project(rectifiedNamedExprs, logicalPlan)
+    Project(namedExprs, logicalPlan)
   }
 
   /**
@@ -4266,7 +4262,7 @@ class Dataset[T] private[sql](
   }
 
   private def convertToUnresolvedIfNeeded(attr: AttributeReference, inputSet: AttributeSet):
-  NamedExpression = {
+    NamedExpression = {
     val attrDatasetIdOpt = if (attr.metadata.contains(Dataset.DATASET_ID_KEY)) {
       Option(attr.metadata.getLong(DATASET_ID_KEY))
     } else {
@@ -4286,7 +4282,6 @@ class Dataset[T] private[sql](
     }
   }
 
-
   ////////////////////////////////////////////////////////////////////////////
   // For Python API
   ////////////////////////////////////////////////////////////////////////////

From ab14974593c65373fb4dee323b767d00b1655ee9 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 14 Mar 2024 15:46:44 -0700
Subject: [PATCH 14/22] SPARK-47320: fixed test failure

---
 sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 1feff155b8706..dc3acb8fc8a4a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -4268,7 +4268,8 @@ class Dataset[T] private[sql](
     } else {
       None
     }
-    if (attrDatasetIdOpt.isDefined && inputSet.find(_ == attr).forall(x => {
+    if (attrDatasetIdOpt.isDefined && inputSet.find(_.canonicalized == attr.canonicalized).
+      forall(x => {
       val datasetIdX = if (x.metadata.contains(DATASET_ID_KEY)) {
         x.metadata.getLong(DATASET_ID_KEY)
       } else {

From bf1cd922dbd31c1bda5a06ddc052da1b38630635 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 14 Mar 2024 19:19:35 -0700
Subject: [PATCH 15/22] SPARK-47320: fixed test failure

---
 .../scala/org/apache/spark/sql/Dataset.scala  | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index dc3acb8fc8a4a..e54211d7e7321 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1182,9 +1182,19 @@ class Dataset[T] private[sql](
       Join(logicalPlan, right.logicalPlan,
         JoinType(joinType), None, JoinHint.NONE))
       .queryExecution.analyzed.asInstanceOf[Join]
-    val inputSet = planPart1.outputSet
+
+    val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
+    val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
     val joinExprsRectified = joinExprs.map(_.expr transformUp {
-      case attr: AttributeReference => convertToUnresolvedIfNeeded(attr, inputSet)
+      case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
+        val attribTagId = attr.metadata.getLong(Dataset.DATASET_ID_KEY)
+        if (!planPart1.outputSet.contains(attr) ||
+          (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
+          (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
+          UnresolvedAttributeWithTag(attr, attribTagId)
+        } else {
+          attr
+        }
     })
     Join(planPart1.left, planPart1.right, JoinType(joinType), joinExprsRectified, JoinHint.NONE)
   }
@@ -1597,9 +1607,10 @@ class Dataset[T] private[sql](
 
       case other => other
     }
-    val inputSet = logicalPlan.outputSet
+    val inputForProj = logicalPlan.output
     val namedExprs = untypedCols.map(ne => (ne.named transformUp {
-      case attr: AttributeReference => convertToUnresolvedIfNeeded(attr, inputSet)
+      case attr: AttributeReference if shouldChangeToUnresolved(attr, inputForProj, id) =>
+        UnresolvedAttributeWithTag(attr, attr.metadata.getLong(DATASET_ID_KEY))
     }).asInstanceOf[NamedExpression])
     Project(namedExprs, logicalPlan)
   }
@@ -4261,26 +4272,24 @@ class Dataset[T] private[sql](
     queryExecution.analyzed.semanticHash()
   }
 
-  private def convertToUnresolvedIfNeeded(attr: AttributeReference, inputSet: AttributeSet):
-    NamedExpression = {
-    val attrDatasetIdOpt = if (attr.metadata.contains(Dataset.DATASET_ID_KEY)) {
+  private def shouldChangeToUnresolved(
+      attr: AttributeReference,
+      input: Seq[Attribute],
+      dataSetIdOfInput: Long): Boolean = {
+    val attrDatasetIdOpt = if (attr.metadata.contains(DATASET_ID_KEY)) {
       Option(attr.metadata.getLong(DATASET_ID_KEY))
     } else {
       None
     }
-    if (attrDatasetIdOpt.isDefined && inputSet.find(_.canonicalized == attr.canonicalized).
+    (attrDatasetIdOpt.isDefined && input.filter(_.canonicalized == attr.canonicalized).
       forall(x => {
       val datasetIdX = if (x.metadata.contains(DATASET_ID_KEY)) {
         x.metadata.getLong(DATASET_ID_KEY)
       } else {
-        Dataset.this.id
+        dataSetIdOfInput
       }
       attrDatasetIdOpt.get != datasetIdX
-    })) {
-      UnresolvedAttributeWithTag(attr, attrDatasetIdOpt.get)
-    } else {
-      attr
-    }
+    }))
   }
 
   ////////////////////////////////////////////////////////////////////////////

From 611847e05506d6a0066f0134050fc02a371e926a Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Thu, 14 Mar 2024 22:34:49 -0700
Subject: [PATCH 16/22] SPARK-47320: refcatored code

---
 .../scala/org/apache/spark/sql/Dataset.scala  | 51 ++++++++++++++-----
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index e54211d7e7321..1817becb85f84 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1182,12 +1182,16 @@ class Dataset[T] private[sql](
       Join(logicalPlan, right.logicalPlan,
         JoinType(joinType), None, JoinHint.NONE))
       .queryExecution.analyzed.asInstanceOf[Join]
+    val inputForCondn = planPart1.output
 
     val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
     val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
-    val joinExprsRectified = joinExprs.map(_.expr transformUp {
+
+
+    val joinExprsRectified1 = joinExprs.map(_.expr transformUp {
       case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
         val attribTagId = attr.metadata.getLong(Dataset.DATASET_ID_KEY)
+
         if (!planPart1.outputSet.contains(attr) ||
           (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
           (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
@@ -1196,6 +1200,20 @@ class Dataset[T] private[sql](
           attr
         }
     })
+
+    val joinExprsRectified = joinExprs.map(_.expr transformUp {
+      case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) =>
+        // For attribute to remain attribute and not to UnResolved, only one leg should be tru
+        val leftLegWrong = isIncorrectlyResolved(attr, planPart1.left.outputSet,
+          leftTagIdMap.getOrElse(HashSet.empty[Long]))
+        val rightLegWrong = isIncorrectlyResolved(attr, planPart1.right.outputSet,
+          rightTagIdMap.getOrElse(HashSet.empty[Long]))
+        if (!planPart1.outputSet.contains(attr) || leftLegWrong || rightLegWrong) {
+          UnresolvedAttributeWithTag(attr, attr.metadata.getLong(DATASET_ID_KEY))
+        } else {
+          attr
+        }
+    })
     Join(planPart1.left, planPart1.right, JoinType(joinType), joinExprsRectified, JoinHint.NONE)
   }
 
@@ -1607,9 +1625,11 @@ class Dataset[T] private[sql](
 
       case other => other
     }
-    val inputForProj = logicalPlan.output
+    val inputForProj = logicalPlan.outputSet
     val namedExprs = untypedCols.map(ne => (ne.named transformUp {
-      case attr: AttributeReference if shouldChangeToUnresolved(attr, inputForProj, id) =>
+      case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) &&
+        (!inputForProj.contains(attr) ||
+          isIncorrectlyResolved(attr, inputForProj, HashSet(id))) =>
         UnresolvedAttributeWithTag(attr, attr.metadata.getLong(DATASET_ID_KEY))
     }).asInstanceOf[NamedExpression])
     Project(namedExprs, logicalPlan)
@@ -4272,24 +4292,29 @@ class Dataset[T] private[sql](
     queryExecution.analyzed.semanticHash()
   }
 
-  private def shouldChangeToUnresolved(
+  private def isIncorrectlyResolved(
       attr: AttributeReference,
-      input: Seq[Attribute],
-      dataSetIdOfInput: Long): Boolean = {
+      input: AttributeSet,
+      dataSetIdOfInput: HashSet[Long]): Boolean = {
     val attrDatasetIdOpt = if (attr.metadata.contains(DATASET_ID_KEY)) {
       Option(attr.metadata.getLong(DATASET_ID_KEY))
     } else {
       None
     }
-    (attrDatasetIdOpt.isDefined && input.filter(_.canonicalized == attr.canonicalized).
-      forall(x => {
-      val datasetIdX = if (x.metadata.contains(DATASET_ID_KEY)) {
-        x.metadata.getLong(DATASET_ID_KEY)
+    attrDatasetIdOpt.forall(attrId => {
+      val matchingInputset = input.filter(_.canonicalized == attr.canonicalized)
+      if (matchingInputset.isEmpty) {
+        true
       } else {
-        dataSetIdOfInput
+        matchingInputset.forall(x => {
+          if (x.metadata.contains(DATASET_ID_KEY)) {
+            attrId != x.metadata.getLong(DATASET_ID_KEY)
+          } else {
+            !dataSetIdOfInput.contains(attrId)
+          }
+        })
       }
-      attrDatasetIdOpt.get != datasetIdX
-    }))
+    })
   }
 
   ////////////////////////////////////////////////////////////////////////////

From 4501ae5d180adbe19698510c98d5f51390b93b5f Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 15 Mar 2024 15:05:16 -0700
Subject: [PATCH 17/22] SPARK-47320: removed dead code

---
 .../scala/org/apache/spark/sql/Dataset.scala   | 18 +-----------------
 1 file changed, 1 insertion(+), 17 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 1817becb85f84..1186cd00ec2f4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1180,27 +1180,11 @@ class Dataset[T] private[sql](
       joinType: String) = {
     val planPart1 = withPlan(
       Join(logicalPlan, right.logicalPlan,
-        JoinType(joinType), None, JoinHint.NONE))
-      .queryExecution.analyzed.asInstanceOf[Join]
-    val inputForCondn = planPart1.output
+        JoinType(joinType), None, JoinHint.NONE)).queryExecution.analyzed.asInstanceOf[Join]
 
     val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
     val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
 
-
-    val joinExprsRectified1 = joinExprs.map(_.expr transformUp {
-      case attr: AttributeReference if attr.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        val attribTagId = attr.metadata.getLong(Dataset.DATASET_ID_KEY)
-
-        if (!planPart1.outputSet.contains(attr) ||
-          (planPart1.left.outputSet.contains(attr) && !leftTagIdMap.contains(attribTagId)) ||
-          (planPart1.right.outputSet.contains(attr) && !rightTagIdMap.contains(attribTagId))) {
-          UnresolvedAttributeWithTag(attr, attribTagId)
-        } else {
-          attr
-        }
-    })
-
     val joinExprsRectified = joinExprs.map(_.expr transformUp {
       case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) =>
         // For attribute to remain attribute and not to UnResolved, only one leg should be tru

From 3b8383d6ecedeca740cff646e6ecf2ae8c26973e Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 15 Mar 2024 21:02:05 -0700
Subject: [PATCH 18/22] SPARK-47320: refactored the code to remove
 UnresolvedAttributeWithTag, instead marking the UnresolvedAttribute using a
 tag

---
 .../analysis/ColumnResolutionHelper.scala     | 150 ++++++++++--------
 .../sql/catalyst/analysis/unresolved.scala    |   4 +-
 .../catalyst/plans/logical/LogicalPlan.scala  |   1 +
 .../scala/org/apache/spark/sql/Dataset.scala  |  22 ++-
 4 files changed, 102 insertions(+), 75 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index 4aac2c6c70676..44ac31f32da87 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -134,7 +134,6 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       expr: Expression,
       resolveColumnByName: Seq[String] => Option[Expression],
       getAttrCandidates: () => Seq[Attribute],
-      resolveOnDatasetId: (Long, String) => Option[NamedExpression],
       throws: Boolean,
       includeLastResort: Boolean): Expression = {
     def innerResolve(e: Expression, isTopLevel: Boolean): Expression = withOrigin(e.origin) {
@@ -157,9 +156,6 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
           }
           matched(ordinal)
 
-        case u @ UnresolvedAttributeWithTag(attr, id) =>
-          resolveOnDatasetId(id, attr.name).getOrElse(attr)
-
         case u @ UnresolvedAttribute(nameParts) =>
           val result = withPosition(u) {
             resolveColumnByName(nameParts).orElse(resolveLiteralFunction(nameParts)).map {
@@ -456,7 +452,6 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
         plan.resolve(nameParts, conf.resolver)
       },
       getAttrCandidates = () => plan.output,
-      resolveOnDatasetId = (_, _) => None,
       throws = throws,
       includeLastResort = includeLastResort)
   }
@@ -482,57 +477,6 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
         assert(q.children.length == 1)
         q.children.head.output
       },
-
-      resolveOnDatasetId = (datasetid: Long, name: String) => {
-        def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[(LogicalPlan, Int)] = {
-          var currentLp = lp
-          var depth = 0
-          while(true) {
-            if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetid))) {
-              return Option(currentLp, depth)
-            } else {
-              if (currentLp.children.size == 1) {
-                currentLp = currentLp.children.head
-              } else {
-                // leaf node or node is a binary node
-                return None
-              }
-            }
-            depth += 1
-          }
-          None
-        }
-
-        val binaryNodeOpt = q.collectFirst {
-          case bn: BinaryNode => bn
-        }
-
-        val resolveOnAttribs = binaryNodeOpt match {
-          case Some(bn) =>
-            val leftDefOpt = findUnaryNodeMatchingTagId(bn.left)
-            val rightDefOpt = findUnaryNodeMatchingTagId(bn.right)
-            (leftDefOpt, rightDefOpt) match {
-
-              case (None, Some((lp, _))) => lp.output
-
-              case (Some((lp, _)), None) => lp.output
-
-              case  (Some((lp1, depth1)), Some((lp2, depth2))) =>
-                if (depth1 == depth2) {
-                  q.children.head.output
-                } else if (depth1 < depth2) {
-                  lp1.output
-                } else {
-                  lp2.output
-                }
-
-              case _ => q.children.head.output
-            }
-
-          case _ => q.children.head.output
-        }
-        AttributeSeq.fromNormalOutput(resolveOnAttribs).resolve(Seq(name), conf.resolver)
-      },
       throws = true,
       includeLastResort = includeLastResort)
   }
@@ -574,24 +518,90 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
     case _ => e
   }
 
+  private def resolveUsingDatasetId(
+      ua: UnresolvedAttribute,
+      left: LogicalPlan,
+      right: LogicalPlan,
+      datasetId: Long): Option[NamedExpression] = {
+    def findUnaryNodeMatchingTagId(lp: LogicalPlan): Option[(LogicalPlan, Int)] = {
+      var currentLp = lp
+      var depth = 0
+      while (true) {
+        if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetId))) {
+          return Option(currentLp, depth)
+        } else {
+          if (currentLp.children.size == 1) {
+            currentLp = currentLp.children.head
+          } else {
+            // leaf node or node is a binary node
+            return None
+          }
+        }
+        depth += 1
+      }
+      None
+    }
+
+    val leftDefOpt = findUnaryNodeMatchingTagId(left)
+    val rightDefOpt = findUnaryNodeMatchingTagId(right)
+    val resolveOnAttribs = (leftDefOpt, rightDefOpt) match {
+
+      case (None, Some((lp, _))) => lp.output
+
+      case (Some((lp, _)), None) => lp.output
+
+      case (Some((lp1, depth1)), Some((lp2, depth2))) =>
+        if (depth1 == depth2) {
+          lp1.output
+        } else if (depth1 < depth2) {
+          lp1.output
+        } else {
+          lp2.output
+        }
+
+      case _ => Seq.empty
+    }
+
+    AttributeSeq.fromNormalOutput(resolveOnAttribs).resolve(Seq(ua.name), conf.resolver)
+  }
+
   private def resolveDataFrameColumn(
       u: UnresolvedAttribute,
       q: Seq[LogicalPlan]): Option[NamedExpression] = {
-    val planIdOpt = u.getTagValue(LogicalPlan.PLAN_ID_TAG)
-    if (planIdOpt.isEmpty) return None
-    val planId = planIdOpt.get
-    logDebug(s"Extract plan_id $planId from $u")
-
-    val isMetadataAccess = u.getTagValue(LogicalPlan.IS_METADATA_COL).nonEmpty
-    val (resolved, matched) = resolveDataFrameColumnByPlanId(u, planId, isMetadataAccess, q)
-    if (!matched) {
-      // Can not find the target plan node with plan id, e.g.
-      //  df1 = spark.createDataFrame([Row(a = 1, b = 2, c = 3)]])
-      //  df2 = spark.createDataFrame([Row(a = 1, b = 2)]])
-      //  df1.select(df2.a)   <-   illegal reference df2.a
-      throw QueryCompilationErrors.cannotResolveDataFrameColumn(u)
+
+    val attrWithDatasetIdOpt = u.getTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG)
+    val resolvedOpt = if (attrWithDatasetIdOpt.isDefined) {
+      val did = attrWithDatasetIdOpt.get
+      if (q.size == 1) {
+        val binaryNodeOpt = q.head.collectFirst {
+          case bn: BinaryNode => bn
+        }
+        binaryNodeOpt.flatMap(bn => resolveUsingDatasetId(u, bn.left, bn.right, did))
+      } else if (q.size == 2) {
+        resolveUsingDatasetId(u, q(0), q(1), did)
+      } else {
+        None
+      }
+    } else {
+      val planIdOpt = u.getTagValue(LogicalPlan.PLAN_ID_TAG)
+      if (planIdOpt.isEmpty) {
+        None
+      } else {
+        val planId = planIdOpt.get
+        logDebug(s"Extract plan_id $planId from $u")
+        val isMetadataAccess = u.getTagValue(LogicalPlan.IS_METADATA_COL).nonEmpty
+        val (resolved, matched) = resolveDataFrameColumnByPlanId(u, planId, isMetadataAccess, q)
+        if (!matched) {
+          // Can not find the target plan node with plan id, e.g.
+          //  df1 = spark.createDataFrame([Row(a = 1, b = 2, c = 3)]])
+          //  df2 = spark.createDataFrame([Row(a = 1, b = 2)]])
+          //  df1.select(df2.a)   <-   illegal reference df2.a
+          throw QueryCompilationErrors.cannotResolveDataFrameColumn(u)
+        }
+        resolved
+      }
     }
-    resolved
+    resolvedOpt
   }
 
   private def resolveDataFrameColumnByPlanId(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index 397351e0c1fdd..a46105dff8538 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -267,7 +267,7 @@ case class UnresolvedAttribute(nameParts: Seq[String]) extends Attribute with Un
     nameParts.length == 1 && nameParts.head.equalsIgnoreCase(token)
   }
 }
-
+/*
 case class UnresolvedAttributeWithTag(attribute: Attribute, datasetId: Long) extends Attribute with
   Unevaluable {
   def name: String = attribute.name
@@ -309,6 +309,8 @@ case class UnresolvedAttributeWithTag(attribute: Attribute, datasetId: Long) ext
   def equalsIgnoreCase(token: String): Boolean = token.equalsIgnoreCase(attribute.name)
 }
 
+ */
+
 object UnresolvedAttribute extends AttributeNameParser {
   /**
    * Creates an [[UnresolvedAttribute]], parsing segments separated by dots ('.').
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
index a9b130c981ac0..45cce59668209 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -201,6 +201,7 @@ object LogicalPlan {
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
   private[spark] val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
+  private[spark] val ATTRIBUTE_DATASET_ID_TAG = TreeNodeTag[Long]("dataset_id")
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 1186cd00ec2f4..1f1edfb0e0a14 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1193,7 +1193,11 @@ class Dataset[T] private[sql](
         val rightLegWrong = isIncorrectlyResolved(attr, planPart1.right.outputSet,
           rightTagIdMap.getOrElse(HashSet.empty[Long]))
         if (!planPart1.outputSet.contains(attr) || leftLegWrong || rightLegWrong) {
-          UnresolvedAttributeWithTag(attr, attr.metadata.getLong(DATASET_ID_KEY))
+          val ua = UnresolvedAttribute(attr.name)
+          ua.copyTagsFrom(attr)
+          ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG,
+            attr.metadata.getLong(DATASET_ID_KEY))
+          ua
         } else {
           attr
         }
@@ -1337,7 +1341,10 @@ class Dataset[T] private[sql](
         joined.left.output(index)
 
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
+        val ua = UnresolvedAttribute(a.name)
+        ua.copyTagsFrom(a)
+        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
+        ua
     }
     val rightAsOfExpr = rightAsOf.expr.transformUp {
       case a: AttributeReference if other.logicalPlan.outputSet.contains(a) =>
@@ -1345,7 +1352,10 @@ class Dataset[T] private[sql](
         joined.right.output(index)
 
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        UnresolvedAttributeWithTag(a, a.metadata.getLong(Dataset.DATASET_ID_KEY))
+        val ua = UnresolvedAttribute(a.name)
+        ua.copyTagsFrom(a)
+        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
+        ua
     }
     withPlan {
       AsOfJoin(
@@ -1614,7 +1624,11 @@ class Dataset[T] private[sql](
       case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) &&
         (!inputForProj.contains(attr) ||
           isIncorrectlyResolved(attr, inputForProj, HashSet(id))) =>
-        UnresolvedAttributeWithTag(attr, attr.metadata.getLong(DATASET_ID_KEY))
+        val ua = UnresolvedAttribute(attr.name)
+        ua.copyTagsFrom(attr)
+        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, attr.metadata.getLong(DATASET_ID_KEY))
+        ua
+
     }).asInstanceOf[NamedExpression])
     Project(namedExprs, logicalPlan)
   }

From f78eaaa61eebaa8ce3da4e9873f262835a022d7f Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 15 Mar 2024 21:07:05 -0700
Subject: [PATCH 19/22] SPARK-47320: removed dead code

---
 .../sql/catalyst/analysis/unresolved.scala    | 43 -------------------
 1 file changed, 43 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index a46105dff8538..7a3cc4bc8e83e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -267,49 +267,6 @@ case class UnresolvedAttribute(nameParts: Seq[String]) extends Attribute with Un
     nameParts.length == 1 && nameParts.head.equalsIgnoreCase(token)
   }
 }
-/*
-case class UnresolvedAttributeWithTag(attribute: Attribute, datasetId: Long) extends Attribute with
-  Unevaluable {
-  def name: String = attribute.name
-
-  override def exprId: ExprId = throw new UnresolvedException("exprId")
-
-  override def dataType: DataType = throw new UnresolvedException("dataType")
-
-  override def nullable: Boolean = throw new UnresolvedException("nullable")
-
-  override def qualifier: Seq[String] = throw new UnresolvedException("qualifier")
-
-  override lazy val resolved = false
-
-  override def newInstance(): UnresolvedAttributeWithTag = this
-
-  override def withNullability(newNullability: Boolean): UnresolvedAttributeWithTag = this
-
-  override def withQualifier(newQualifier: Seq[String]): UnresolvedAttributeWithTag = this
-
-  override def withName(newName: String): UnresolvedAttributeWithTag = this
-
-  override def withMetadata(newMetadata: Metadata): Attribute = this
-
-  override def withExprId(newExprId: ExprId): UnresolvedAttributeWithTag = this
-
-  override def withDataType(newType: DataType): Attribute = this
-
-  final override val nodePatterns: Seq[TreePattern] = Seq(UNRESOLVED_ATTRIBUTE)
-
-  override def toString: String = s"'$name"
-
-  override def sql: String = attribute.sql
-
-  /**
-   * Returns true if this matches the token. This requires the attribute to only have one part in
-   * its name and that matches the given token in a case insensitive way.
-   */
-  def equalsIgnoreCase(token: String): Boolean = token.equalsIgnoreCase(attribute.name)
-}
-
- */
 
 object UnresolvedAttribute extends AttributeNameParser {
   /**

From 11bc2316249df02c887cdbdb9e026f5c9a6e1231 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 15 Mar 2024 23:33:34 -0700
Subject: [PATCH 20/22] SPARK-47320: fixed pyspark failures

---
 .../src/main/scala/org/apache/spark/sql/Dataset.scala     | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 1f1edfb0e0a14..ffb5ae11e9660 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -1193,7 +1193,7 @@ class Dataset[T] private[sql](
         val rightLegWrong = isIncorrectlyResolved(attr, planPart1.right.outputSet,
           rightTagIdMap.getOrElse(HashSet.empty[Long]))
         if (!planPart1.outputSet.contains(attr) || leftLegWrong || rightLegWrong) {
-          val ua = UnresolvedAttribute(attr.name)
+          val ua = UnresolvedAttribute(Seq(attr.name))
           ua.copyTagsFrom(attr)
           ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG,
             attr.metadata.getLong(DATASET_ID_KEY))
@@ -1341,7 +1341,7 @@ class Dataset[T] private[sql](
         joined.left.output(index)
 
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        val ua = UnresolvedAttribute(a.name)
+        val ua = UnresolvedAttribute(Seq(a.name))
         ua.copyTagsFrom(a)
         ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
         ua
@@ -1352,7 +1352,7 @@ class Dataset[T] private[sql](
         joined.right.output(index)
 
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
-        val ua = UnresolvedAttribute(a.name)
+        val ua = UnresolvedAttribute(Seq(a.name))
         ua.copyTagsFrom(a)
         ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
         ua
@@ -1624,7 +1624,7 @@ class Dataset[T] private[sql](
       case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) &&
         (!inputForProj.contains(attr) ||
           isIncorrectlyResolved(attr, inputForProj, HashSet(id))) =>
-        val ua = UnresolvedAttribute(attr.name)
+        val ua = UnresolvedAttribute(Seq(attr.name))
         ua.copyTagsFrom(attr)
         ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, attr.metadata.getLong(DATASET_ID_KEY))
         ua

From 03149d53fbdae034b3f1e61daf0ed5bf2ccc0cb0 Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 29 Mar 2024 15:12:38 -0700
Subject: [PATCH 21/22] SPARK-47320. Modified the code to ensure that for
 unambiguous attributes resolved using datasetId for top level join, the
 behaviour remains unchanged independent of the flag
 spark.sql.analyzer.failAmbiguousSelfJoin value

---
 .../analysis/ColumnResolutionHelper.scala     |  75 ++-
 .../catalyst/plans/logical/LogicalPlan.scala  |   6 +-
 .../scala/org/apache/spark/sql/Dataset.scala  |  30 +-
 .../spark/sql/DataFrameSelfJoinSuite.scala    | 586 ++++++++++--------
 4 files changed, 388 insertions(+), 309 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index 44ac31f32da87..abc59a69bfed9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.util.toPrettySQL
 import org.apache.spark.sql.connector.catalog.{CatalogManager, Identifier}
 import org.apache.spark.sql.errors.{DataTypeErrorsBase, QueryCompilationErrors}
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.MetadataBuilder
 
 trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
 
@@ -518,6 +519,15 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
     case _ => e
   }
 
+  private def stripColumnReferenceMetadata(a: AttributeReference): AttributeReference = {
+    val metadataWithoutId = new MetadataBuilder()
+      .withMetadata(a.metadata)
+      .remove(LogicalPlan.DATASET_ID_KEY)
+      .remove(LogicalPlan.COL_POS_KEY)
+      .build()
+    a.withMetadata(metadataWithoutId)
+  }
+
   private def resolveUsingDatasetId(
       ua: UnresolvedAttribute,
       left: LogicalPlan,
@@ -527,7 +537,8 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       var currentLp = lp
       var depth = 0
       while (true) {
-        if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(_.contains(datasetId))) {
+        if (currentLp.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG).exists(
+          _.contains(datasetId))) {
           return Option(currentLp, depth)
         } else {
           if (currentLp.children.size == 1) {
@@ -550,39 +561,61 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
 
       case (Some((lp, _)), None) => lp.output
 
-      case (Some((lp1, depth1)), Some((lp2, depth2))) =>
-        if (depth1 == depth2) {
-          lp1.output
-        } else if (depth1 < depth2) {
-          lp1.output
-        } else {
-          lp2.output
-        }
+      case (Some((lp1, depth1)), Some((lp2, depth2))) => if (depth1 == depth2) {
+        Seq.empty
+      } else if (depth1 < depth2) {
+        lp1.output
+      } else {
+        lp2.output
+      }
 
       case _ => Seq.empty
     }
-
-    AttributeSeq.fromNormalOutput(resolveOnAttribs).resolve(Seq(ua.name), conf.resolver)
+    if (resolveOnAttribs.isEmpty) {
+      None
+    } else {
+      AttributeSeq.fromNormalOutput(resolveOnAttribs).resolve(Seq(ua.name), conf.resolver)
+    }
   }
 
   private def resolveDataFrameColumn(
       u: UnresolvedAttribute,
       q: Seq[LogicalPlan]): Option[NamedExpression] = {
 
-    val attrWithDatasetIdOpt = u.getTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG)
-    val resolvedOpt = if (attrWithDatasetIdOpt.isDefined) {
-      val did = attrWithDatasetIdOpt.get
-      if (q.size == 1) {
-        val binaryNodeOpt = q.head.collectFirst {
-          case bn: BinaryNode => bn
+    val origAttrOpt = u.getTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG)
+    val resolvedOptWithDatasetId = if (origAttrOpt.isDefined) {
+      val md = origAttrOpt.get.metadata
+      if (md.contains(LogicalPlan.DATASET_ID_KEY)) {
+        val did = md.getLong(LogicalPlan.DATASET_ID_KEY)
+        val resolved = if (q.size == 1) {
+          val binaryNodeOpt = q.head.collectFirst {
+            case bn: BinaryNode => bn
+          }
+          binaryNodeOpt.flatMap(bn => resolveUsingDatasetId(u, bn.left, bn.right, did))
+        } else if (q.size == 2) {
+          resolveUsingDatasetId(u, q(0), q(1), did)
+        } else {
+          None
+        }
+        if (resolved.isEmpty) {
+          if (conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED)) {
+            origAttrOpt
+          } else {
+            origAttrOpt.map(stripColumnReferenceMetadata)
+          }
+        } else {
+          resolved
         }
-        binaryNodeOpt.flatMap(bn => resolveUsingDatasetId(u, bn.left, bn.right, did))
-      } else if (q.size == 2) {
-        resolveUsingDatasetId(u, q(0), q(1), did)
       } else {
-        None
+        origAttrOpt
       }
     } else {
+      None
+    }
+    val resolvedOpt = if (resolvedOptWithDatasetId.isDefined) {
+      resolvedOptWithDatasetId
+    }
+    else {
       val planIdOpt = u.getTagValue(LogicalPlan.PLAN_ID_TAG)
       if (planIdOpt.isEmpty) {
         None
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
index 45cce59668209..1ac3fda25979b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -200,8 +200,10 @@ object LogicalPlan {
   //    to the old code path.
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
-  private[spark] val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
-  private[spark] val ATTRIBUTE_DATASET_ID_TAG = TreeNodeTag[Long]("dataset_id")
+  private[spark] val DATASET_RESOLUTION_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
+  private[spark] val UNRESOLVED_ATTRIBUTE_MD_TAG = TreeNodeTag[AttributeReference]("orig-attr")
+  private[spark] val DATASET_ID_KEY = "__dataset_id"
+  private[spark] val COL_POS_KEY = "__col_position"
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index c01cfc0a92bdc..4b7086432be2f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql
 import java.io.{ByteArrayOutputStream, CharArrayWriter, DataOutputStream}
 
 import scala.annotation.varargs
+import scala.collection.mutable
 import scala.collection.mutable.{ArrayBuffer, HashSet}
 import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
@@ -48,7 +49,7 @@ import org.apache.spark.sql.catalyst.json.{JacksonGenerator, JSONOptions}
 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.trees.TreePattern
+import org.apache.spark.sql.catalyst.trees.{TreeNodeTag, TreePattern}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, IntervalUtils}
 import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId
@@ -72,9 +73,9 @@ 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 = LogicalPlan.DATASET_ID_TAG
+  val DATASET_ID_KEY = LogicalPlan.DATASET_ID_KEY
+  val COL_POS_KEY = LogicalPlan.COL_POS_KEY
+  val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
 
   def apply[T: Encoder](sparkSession: SparkSession, logicalPlan: LogicalPlan): Dataset[T] = {
     val dataset = new Dataset(sparkSession, logicalPlan, implicitly[Encoder[T]])
@@ -228,6 +229,9 @@ class Dataset[T] private[sql](
       dsIds.add(id)
       plan.setTagValue(Dataset.DATASET_ID_TAG, dsIds)
     }
+    val dsIds = plan.getTagValue(Dataset.DATASET_ID_TAG).getOrElse(new HashSet[Long])
+    dsIds.add(id)
+    plan.setTagValue(LogicalPlan.DATASET_RESOLUTION_TAG, dsIds)
     plan
   }
 
@@ -1177,8 +1181,8 @@ class Dataset[T] private[sql](
       Join(logicalPlan, right.logicalPlan,
         JoinType(joinType), None, JoinHint.NONE)).queryExecution.analyzed.asInstanceOf[Join]
 
-    val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_ID_TAG)
-    val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_ID_TAG)
+    val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG)
+    val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG)
 
     val joinExprsRectified = joinExprs.map(_.expr transformUp {
       case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) =>
@@ -1190,8 +1194,7 @@ class Dataset[T] private[sql](
         if (!planPart1.outputSet.contains(attr) || leftLegWrong || rightLegWrong) {
           val ua = UnresolvedAttribute(Seq(attr.name))
           ua.copyTagsFrom(attr)
-          ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG,
-            attr.metadata.getLong(DATASET_ID_KEY))
+          ua.setTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG, attr)
           ua
         } else {
           attr
@@ -1340,7 +1343,7 @@ class Dataset[T] private[sql](
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
         val ua = UnresolvedAttribute(Seq(a.name))
         ua.copyTagsFrom(a)
-        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
+        ua.setTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG, a)
         ua
     }
     val rightAsOfExpr = rightAsOf.expr.transformUp {
@@ -1351,7 +1354,7 @@ class Dataset[T] private[sql](
       case a: AttributeReference if a.metadata.contains(Dataset.DATASET_ID_KEY) =>
         val ua = UnresolvedAttribute(Seq(a.name))
         ua.copyTagsFrom(a)
-        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, a.metadata.getLong(DATASET_ID_KEY))
+        ua.setTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG, a)
         ua
     }
     withPlan {
@@ -1525,8 +1528,8 @@ class Dataset[T] private[sql](
   // `DetectAmbiguousSelfJoin` will remove it.
   private def addDataFrameIdToCol(expr: NamedExpression): NamedExpression = {
     val newExpr = expr transform {
-      case a: AttributeReference
-        if sparkSession.conf.get(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED) =>
+      case a: AttributeReference =>
+        // if sparkSession.conf.get(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED) =>
         val metadata = new MetadataBuilder()
           .withMetadata(a.metadata)
           .putLong(Dataset.DATASET_ID_KEY, id)
@@ -1623,9 +1626,8 @@ class Dataset[T] private[sql](
           isIncorrectlyResolved(attr, inputForProj, HashSet(id))) =>
         val ua = UnresolvedAttribute(Seq(attr.name))
         ua.copyTagsFrom(attr)
-        ua.setTagValue(LogicalPlan.ATTRIBUTE_DATASET_ID_TAG, attr.metadata.getLong(DATASET_ID_KEY))
+        ua.setTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG, attr)
         ua
-
     }).asInstanceOf[NamedExpression])
     Project(namedExprs, logicalPlan)
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
index 32a6deb436f9f..b7fb57b325958 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala
@@ -119,80 +119,78 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   }
 
   test("SPARK-28344: NOT AN ambiguous self join - column ref in join condition") {
-    val df1 = spark.range(3)
-    val df2 = df1.filter($"id" > 0)
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "false",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      // `df1("id") > df2("id")` is always false.
-      checkAnswer(df1.join(df2, df1("id") > df2("id")), Nil)
-
-      // Alias the dataframe and use qualified column names can fix ambiguous self-join.
-      val aliasedDf1 = df1.alias("left")
-      val aliasedDf2 = df2.as("right")
-      checkAnswer(
-        aliasedDf1.join(aliasedDf2, $"left.id" > $"right.id"),
-        Seq(Row(2, 1)))
-    }
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString,
+        SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+        val df1 = spark.range(3)
+        val df2 = df1.filter($"id" > 0)
+        // `df1("id") > df2("id")` is always false.
+        checkAnswer(df1.join(df2, df1("id") > df2("id")), Seq(Row(2, 1)))
         assertCorrectResolution(df1.join(df2, df1("id") > df2("id")),
           Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    }
+
+        // Alias the dataframe and use qualified column names to eliminate all possibilities
+        // of ambiguity in self-join.
+        val aliasedDf1 = df1.alias("left")
+        val aliasedDf2 = df2.as("right")
+        checkAnswer(
+          aliasedDf1.join(aliasedDf2, $"left.id" > $"right.id"),
+          Seq(Row(2, 1)))
+        assertCorrectResolution(aliasedDf1.join(aliasedDf2, $"left.id" > $"right.id"),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+      }
+    })
   }
 
   test("SPARK-28344: Not AN ambiguous self join - Dataset.colRegex as column ref") {
-    val df1 = spark.range(3)
-    val df2 = df1.filter($"id" > 0)
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertCorrectResolution(df1.join(df2, df1.colRegex("id") > df2.colRegex("id")),
-        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    }
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString,
+        SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+        val df1 = spark.range(3)
+        val df2 = df1.filter($"id" > 0)
+        assertCorrectResolution(df1.join(df2, df1.colRegex("id") > df2.colRegex("id")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+      }
+    })
   }
 
   test("SPARK-28344: Not An ambiguous self join - Dataset.col with nested field") {
-    val df1 = spark.read.json(Seq("""{"a": {"b": 1, "c": 1}}""").toDS())
-    val df2 = df1.filter($"a.b" > 0)
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertCorrectResolution( df1.join(df2, df1("a.b") > df2("a.c")),
-        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    }
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString,
+        SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+        val df1 = spark.read.json(Seq("""{"a": {"b": 1, "c": 1}}""").toDS())
+        val df2 = df1.filter($"a.b" > 0)
+        assertCorrectResolution(df1.join(df2, df1("a.b") > df2("a.c")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+      }
+    })
   }
 
-  test("SPARK-28344: fail ambiguous self join - column ref in Project") {
-    val df1 = spark.range(3)
-    val df2 = df1.filter($"id" > 0)
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "false",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      // `df2("id")` actually points to the column of `df1`.
-      checkAnswer(df1.join(df2).select(df2("id")), Seq(0, 0, 1, 1, 2, 2).map(Row(_)))
-
-      // Alias the dataframe and use qualified column names can fix ambiguous self-join.
-      val aliasedDf1 = df1.alias("left")
-      val aliasedDf2 = df2.as("right")
-      checkAnswer(
-        aliasedDf1.join(aliasedDf2).select($"right.id"),
-        Seq(1, 1, 1, 2, 2, 2).map(Row(_)))
-    }
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      val proj1 = df1.join(df2).select(df2("id")).queryExecution.analyzed.asInstanceOf[Project]
-      val join1 = proj1.child.asInstanceOf[Join]
-      assert(proj1.projectList(0).references.subsetOf(join1.right.outputSet))
-    }
+  test("SPARK-28344: Not an ambiguous  - column ref in Project") {
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString,
+        SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+        val df1 = spark.range(3)
+        val df2 = df1.filter($"id" > 0)
+        // `df2("id")` actually points to the column of `df1`.
+        checkAnswer(df1.join(df2).select(df2("id")), Seq(1, 1, 1, 2, 2, 2).map(Row(_)))
+
+        // Alias the dataframe and use qualified column names can fix ambiguous self-join.
+        val aliasedDf1 = df1.alias("left")
+        val aliasedDf2 = df2.as("right")
+        checkAnswer(
+          aliasedDf1.join(aliasedDf2).select($"right.id"),
+          Seq(1, 1, 1, 2, 2, 2).map(Row(_)))
+
+        val proj1 = df1.join(df2).select(df2("id")).queryExecution.analyzed.asInstanceOf[Project]
+        val join1 = proj1.child.asInstanceOf[Join]
+        assert(proj1.projectList(0).references.subsetOf(join1.right.outputSet))
+      }
+    })
   }
 
   test("SPARK-28344: fail ambiguous self join - join three tables") {
@@ -204,12 +202,13 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "false",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      // `df2("id") < df3("id")` is always false
-      checkAnswer(df1.join(df2).join(df3, df2("id") < df3("id")), Nil)
+      // Here df3("id") is unambiguous, df2("id") is ambiguous. default resolves to df1
+      checkAnswer(df1.join(df2).join(df3, df2("id") < df3("id")),
+        Seq(Row(0, 1, 1), Row(0, 2, 1), Row(0, 1, 2), Row(0, 2, 2), Row(1, 1, 2), Row(1, 2, 2)))
       // `df2("id")` actually points to the column of `df1`.
-      checkAnswer(
+    checkAnswer(
         df1.join(df4).join(df2).select(df2("id")),
-        Seq(0, 0, 1, 1, 2, 2).map(Row(_)))
+        Seq(1, 2, 1, 2, 1, 2).map(Row(_)))
       // `df4("id")` is not ambiguous.
       checkAnswer(
         df1.join(df4).join(df2).select(df4("id")),
@@ -240,19 +239,21 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   }
 
   test("SPARK-28344: don't fail if there is no ambiguous self join") {
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true") {
-      val df = Seq(1, 1, 2, 2).toDF("a")
-      val w = Window.partitionBy(df("a"))
-      checkAnswer(
-        df.select(df("a").alias("x"), sum(df("a")).over(w)),
-        Seq((1, 2), (1, 2), (2, 4), (2, 4)).map(Row.fromTuple))
-
-      val joined = df.join(spark.range(1)).select($"a")
-      checkAnswer(
-        joined.select(joined("a").alias("x"), sum(joined("a")).over(w)),
-        Seq((1, 2), (1, 2), (2, 4), (2, 4)).map(Row.fromTuple))
-    }
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString) {
+        val df = Seq(1, 1, 2, 2).toDF("a")
+        val w = Window.partitionBy(df("a"))
+        checkAnswer(
+          df.select(df("a").alias("x"), sum(df("a")).over(w)),
+          Seq((1, 2), (1, 2), (2, 4), (2, 4)).map(Row.fromTuple))
+
+        val joined = df.join(spark.range(1)).select($"a")
+        checkAnswer(
+          joined.select(joined("a").alias("x"), sum(joined("a")).over(w)),
+          Seq((1, 2), (1, 2), (2, 4), (2, 4)).map(Row.fromTuple))
+      }
+    })
   }
 
   test("SPARK-33071/SPARK-33536: Avoid changing dataset_id of LogicalPlan in join() " +
@@ -333,29 +334,51 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   }
 
   test("SPARK-35454: Not an ambiguous self join - toDF") {
-    val df1 = spark.range(3).toDF()
-    val df2 = df1.filter($"id" > 0).toDF()
-
-    withSQLConf(
-      SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "true",
-      SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      assertCorrectResolution(df1.join(df2, df1.col("id") > df2.col("id")),
-        Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    }
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString,
+        SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+        val df1 = spark.range(3).toDF()
+        val df2 = df1.filter($"id" > 0).toDF()
+        assertCorrectResolution(df1.join(df2, df1.col("id") > df2.col("id")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+      }
+    })
   }
 
   test("SPARK-35454: fail ambiguous self join - join four tables") {
     val df1 = spark.range(3).select($"id".as("a"), $"id".as("b"))
     val df2 = df1.filter($"a" > 0).select("b")
     val df3 = df1.filter($"a" <= 2).select("b")
-    val df4 = df1.filter($"b" <= 2)
+    val df4 = df1.filter($"b" <= 2).as("temp")
     val df5 = spark.range(1)
 
     withSQLConf(
       SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> "false",
       SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
-      // `df2("b") < df4("b")` is always false
-      checkAnswer(df1.join(df2).join(df3).join(df4, df2("b") < df4("b")), Nil)
+
+      //  df4("b") is unambiguous
+      checkAnswer(df1.join(df2).join(df3).join(df4, df2("b") < df4("b")),
+        Seq(
+          Row(0, 0, 1, 0, 1, 1),
+          Row(0, 0, 1, 1, 1, 1),
+          Row(0, 0, 1, 2, 1, 1),
+          Row(0, 0, 2, 0, 1, 1),
+          Row(0, 0, 2, 1, 1, 1),
+          Row(0, 0, 2, 2, 1, 1),
+          Row(0, 0, 1, 0, 2, 2),
+          Row(0, 0, 1, 1, 2, 2),
+          Row(0, 0, 1, 2, 2, 2),
+          Row(0, 0, 2, 0, 2, 2),
+          Row(0, 0, 2, 1, 2, 2),
+          Row(0, 0, 2, 2, 2, 2),
+          Row(1, 1, 1, 0, 2, 2),
+          Row(1, 1, 1, 1, 2, 2),
+          Row(1, 1, 1, 2, 2, 2),
+          Row(1, 1, 2, 0, 2, 2),
+          Row(1, 1, 2, 1, 2, 2),
+          Row(1, 1, 2, 2, 2, 2)
+        ))
       // `df2("b")` actually points to the column of `df1`.
       checkAnswer(
         df1.join(df2).join(df5).join(df4).select(df2("b")),
@@ -391,142 +414,146 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   test("SPARK-36874: DeduplicateRelations should copy dataset_id tag " +
     "to avoid ambiguous self join") {
     // Test for Project
-
-    val df1 = Seq((1, 2, "A1"), (2, 1, "A2")).toDF("key1", "key2", "value")
-    val df2 = df1.filter($"value" === "A2")
-    assertCorrectResolution(df1.join(df2, df1("key1") === df2("key2")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df2.join(df1, df1("key1") === df2("key2")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString) {
+        val df1 = Seq((1, 2, "A1"), (2, 1, "A2")).toDF("key1", "key2", "value")
+        val df2 = df1.filter($"value" === "A2")
+        assertCorrectResolution(df1.join(df2, df1("key1") === df2("key2")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df2.join(df1, df1("key1") === df2("key2")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
 
-    // Test for SerializeFromObject
-    val df3 = spark.sparkContext.parallelize(1 to 10).map(x => (x, x)).toDF()
-    val df4 = df3.filter($"_1" <=> 0)
-    assertCorrectResolution(df3.join(df4, df3("_1") === df4("_2")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df4.join(df3, df3("_1") === df4("_2")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test for SerializeFromObject
+        val df3 = spark.sparkContext.parallelize(1 to 10).map(x => (x, x)).toDF()
+        val df4 = df3.filter($"_1" <=> 0)
+        assertCorrectResolution(df3.join(df4, df3("_1") === df4("_2")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df4.join(df3, df3("_1") === df4("_2")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
-    // Test For Aggregate
-    val df5 = df1.groupBy($"key1").agg(count($"value") as "count")
-    val df6 = df5.filter($"key1" > 0)
-    assertCorrectResolution(df5.join(df6, df5("key1") === df6("count")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df6.join(df5, df5("key1") === df6("count")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    // Test for MapInPandas
-    val mapInPandasUDF = PythonUDF("mapInPandasUDF", null,
-      StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
-      Seq.empty,
-      PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
-      true)
-    val df7 = df1.mapInPandas(mapInPandasUDF)
-    val df8 = df7.filter($"x" > 0)
-    assertCorrectResolution(df7.join(df8, df7("x") === df8("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df8.join(df7, df7("x") === df8("y")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    // Test for FlatMapGroupsInPandas
-    val flatMapGroupsInPandasUDF = PythonUDF("flagMapGroupsInPandasUDF", null,
-      StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
-      Seq.empty,
-      PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
-      true)
-    val df9 = df1.groupBy($"key1").flatMapGroupsInPandas(flatMapGroupsInPandasUDF)
-    val df10 = df9.filter($"x" > 0)
-    assertCorrectResolution(df9.join(df10, df9("x") === df10("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df10.join(df9, df9("x") === df10("y")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    // Test for FlatMapCoGroupsInPandas
-    val flatMapCoGroupsInPandasUDF = PythonUDF("flagMapCoGroupsInPandasUDF", null,
-      StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
-      Seq.empty,
-      PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
-      true)
-    val df11 = df1.groupBy($"key1").flatMapCoGroupsInPandas(
-      df1.groupBy($"key2"), flatMapCoGroupsInPandasUDF)
-    val df12 = df11.filter($"x" > 0)
-    assertCorrectResolution(df11.join(df12, df11("x") === df12("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df12.join(df11, df11("x") === df12("y")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test For Aggregate
+        val df5 = df1.groupBy($"key1").agg(count($"value") as "count")
+        val df6 = df5.filter($"key1" > 0)
+        assertCorrectResolution(df5.join(df6, df5("key1") === df6("count")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df6.join(df5, df5("key1") === df6("count")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        // Test for MapInPandas
+        val mapInPandasUDF = PythonUDF("mapInPandasUDF", null,
+          StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
+          Seq.empty,
+          PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
+          true)
+        val df7 = df1.mapInPandas(mapInPandasUDF)
+        val df8 = df7.filter($"x" > 0)
+        assertCorrectResolution(df7.join(df8, df7("x") === df8("y")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df8.join(df7, df7("x") === df8("y")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        // Test for FlatMapGroupsInPandas
+        val flatMapGroupsInPandasUDF = PythonUDF("flagMapGroupsInPandasUDF", null,
+          StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
+          Seq.empty,
+          PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+          true)
+        val df9 = df1.groupBy($"key1").flatMapGroupsInPandas(flatMapGroupsInPandasUDF)
+        val df10 = df9.filter($"x" > 0)
+        assertCorrectResolution(df9.join(df10, df9("x") === df10("y")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df10.join(df9, df9("x") === df10("y")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        // Test for FlatMapCoGroupsInPandas
+        val flatMapCoGroupsInPandasUDF = PythonUDF("flagMapCoGroupsInPandasUDF", null,
+          StructType(Seq(StructField("x", LongType), StructField("y", LongType))),
+          Seq.empty,
+          PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+          true)
+        val df11 = df1.groupBy($"key1").flatMapCoGroupsInPandas(
+          df1.groupBy($"key2"), flatMapCoGroupsInPandasUDF)
+        val df12 = df11.filter($"x" > 0)
+        assertCorrectResolution(df11.join(df12, df11("x") === df12("y")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df12.join(df11, df11("x") === df12("y")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
 
-    // Test for AttachDistributedSequence
-    val df13 = df1.withSequenceColumn("seq")
-    val df14 = df13.filter($"value" === "A2")
-    assertCorrectResolution(df13.join(df14, df13("key1") === df14("key2")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df14.join(df13, df13("key1") === df14("key2")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-    // Test for Generate
-    // Ensure that the root of the plan is Generate
-    val df15 = Seq((1, Seq(1, 2, 3))).toDF("a", "intList").select($"a", explode($"intList"))
-      .queryExecution.optimizedPlan.find(_.isInstanceOf[Generate]).get.toDF()
-    val df16 = df15.filter($"a" > 0)
-    assertCorrectResolution(df15.join(df16, df15("a") === df16("col")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df16.join(df15, df15("a") === df16("col")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    // Test for Expand
-    // Ensure that the root of the plan is Expand
-    val df17 =
-      Expand(
-        Seq(Seq($"key1".expr, $"key2".expr)),
-        Seq(
-          AttributeReference("x", IntegerType)(),
-          AttributeReference("y", IntegerType)()),
-        df1.queryExecution.logical).toDF()
-    val df18 = df17.filter($"x" > 0)
-    assertCorrectResolution(df17.join(df18, df17("x") === df18("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df18.join(df17, df17("x") === df18("y")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-    // Test for Window
-    val dfWithTS = spark.sql("SELECT timestamp'2021-10-15 01:52:00' time, 1 a, 2 b")
-    // Ensure that the root of the plan is Window
-    val df19 = WindowPlan(
-      Seq(Alias(dfWithTS("time").expr, "ts")()),
-      Seq(dfWithTS("a").expr),
-      Seq(SortOrder(dfWithTS("a").expr, Ascending)),
-      dfWithTS.queryExecution.logical).toDF()
-    val df20 = df19.filter($"a" > 0)
-    assertCorrectResolution(df19.join(df20, df19("a") === df20("b")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df20.join(df19, df19("a") === df20("b")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-    // Test for ScriptTransformation
-    val ioSchema =
-      ScriptInputOutputSchema(
-        Seq(("TOK_TABLEROWFORMATFIELD", ","),
-          ("TOK_TABLEROWFORMATCOLLITEMS", "#"),
-          ("TOK_TABLEROWFORMATMAPKEYS", "@"),
-          ("TOK_TABLEROWFORMATNULL", "null"),
-          ("TOK_TABLEROWFORMATLINES", "\n")),
-        Seq(("TOK_TABLEROWFORMATFIELD", ","),
-          ("TOK_TABLEROWFORMATCOLLITEMS", "#"),
-          ("TOK_TABLEROWFORMATMAPKEYS", "@"),
-          ("TOK_TABLEROWFORMATNULL", "null"),
-          ("TOK_TABLEROWFORMATLINES", "\n")), None, None,
-        List.empty, List.empty, None, None, false)
-    // Ensure that the root of the plan is ScriptTransformation
-    val df21 = ScriptTransformation(
-      "cat",
-      Seq(
-        AttributeReference("x", IntegerType)(),
-        AttributeReference("y", IntegerType)()),
-      df1.queryExecution.logical,
-      ioSchema).toDF()
-    val df22 = df21.filter($"x" > 0)
-    assertCorrectResolution(df21.join(df22, df21("x") === df22("y")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-    assertCorrectResolution(df22.join(df21, df21("x") === df22("y")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test for AttachDistributedSequence
+        val df13 = df1.withSequenceColumn("seq")
+        val df14 = df13.filter($"value" === "A2")
+        assertCorrectResolution(df13.join(df14, df13("key1") === df14("key2")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df14.join(df13, df13("key1") === df14("key2")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test for Generate
+        // Ensure that the root of the plan is Generate
+        val df15 = Seq((1, Seq(1, 2, 3))).toDF("a", "intList").select($"a", explode($"intList"))
+          .queryExecution.optimizedPlan.find(_.isInstanceOf[Generate]).get.toDF()
+        val df16 = df15.filter($"a" > 0)
+        assertCorrectResolution(df15.join(df16, df15("a") === df16("col")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df16.join(df15, df15("a") === df16("col")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        // Test for Expand
+        // Ensure that the root of the plan is Expand
+        val df17 =
+        Expand(
+          Seq(Seq($"key1".expr, $"key2".expr)),
+          Seq(
+            AttributeReference("x", IntegerType)(),
+            AttributeReference("y", IntegerType)()),
+          df1.queryExecution.logical).toDF()
+        val df18 = df17.filter($"x" > 0)
+        assertCorrectResolution(df17.join(df18, df17("x") === df18("y")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df18.join(df17, df17("x") === df18("y")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test for Window
+        val dfWithTS = spark.sql("SELECT timestamp'2021-10-15 01:52:00' time, 1 a, 2 b")
+        // Ensure that the root of the plan is Window
+        val df19 = WindowPlan(
+          Seq(Alias(dfWithTS("time").expr, "ts")()),
+          Seq(dfWithTS("a").expr),
+          Seq(SortOrder(dfWithTS("a").expr, Ascending)),
+          dfWithTS.queryExecution.logical).toDF()
+        val df20 = df19.filter($"a" > 0)
+        assertCorrectResolution(df19.join(df20, df19("a") === df20("b")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df20.join(df19, df19("a") === df20("b")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+        // Test for ScriptTransformation
+        val ioSchema =
+          ScriptInputOutputSchema(
+            Seq(("TOK_TABLEROWFORMATFIELD", ","),
+              ("TOK_TABLEROWFORMATCOLLITEMS", "#"),
+              ("TOK_TABLEROWFORMATMAPKEYS", "@"),
+              ("TOK_TABLEROWFORMATNULL", "null"),
+              ("TOK_TABLEROWFORMATLINES", "\n")),
+            Seq(("TOK_TABLEROWFORMATFIELD", ","),
+              ("TOK_TABLEROWFORMATCOLLITEMS", "#"),
+              ("TOK_TABLEROWFORMATMAPKEYS", "@"),
+              ("TOK_TABLEROWFORMATNULL", "null"),
+              ("TOK_TABLEROWFORMATLINES", "\n")), None, None,
+            List.empty, List.empty, None, None, false)
+        // Ensure that the root of the plan is ScriptTransformation
+        val df21 = ScriptTransformation(
+          "cat",
+          Seq(
+            AttributeReference("x", IntegerType)(),
+            AttributeReference("y", IntegerType)()),
+          df1.queryExecution.logical,
+          ioSchema).toDF()
+        val df22 = df21.filter($"x" > 0)
+        assertCorrectResolution(df21.join(df22, df21("x") === df22("y")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
+        assertCorrectResolution(df22.join(df21, df21("x") === df22("y")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+      }
+    })
   }
 
   test("SPARK-35937: GetDateFieldOperations should skip unresolved nodes") {
@@ -561,61 +588,76 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession {
   }
 
   test("SPARK-47217: deduplication of project causes ambiguity in resolution") {
-    val df = Seq((1, 2)).toDF("a", "b")
-    val df2 = df.select(df("a").as("aa"), df("b").as("bb"))
-    val df3 = df2.join(df, df2("bb") === df("b")).select(df2("aa"), df("a"))
-    checkAnswer(
-      df3,
-      Row(1, 1) :: Nil)
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString) {
+        val df = Seq((1, 2)).toDF("a", "b")
+        val df2 = df.select(df("a").as("aa"), df("b").as("bb"))
+        val df3 = df2.join(df, df2("bb") === df("b")).select(df2("aa"), df("a"))
+        checkAnswer(
+          df3,
+          Row(1, 1) :: Nil)
+      }
+    })
   }
 
   test("SPARK-47217: deduplication in nested joins with join attribute aliased") {
-    val df1 = Seq((1, 2)).toDF("a", "b")
-    val df2 = Seq((1, 2)).toDF("aa", "bb")
-    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
-      df2("aa"), df1("b"))
-
-    assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-
-    assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    val proj1 = df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
-      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
-    val join1 = proj1.child.asInstanceOf[Join]
-    assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
-    assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString) {
+        val df1 = Seq((1, 2)).toDF("a", "b")
+        val df2 = Seq((1, 2)).toDF("aa", "bb")
+        val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a").as("aaa"),
+          df2("aa"), df1("b"))
+
+        assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
 
-    val proj2 = df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
-      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
-    val join2 = proj2.child.asInstanceOf[Join]
-    assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
-    assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
+        assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        val proj1 = df1Joindf2.join(df1, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
+          df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+        val join1 = proj1.child.asInstanceOf[Join]
+        assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
+        assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+
+        val proj2 = df1.join(df1Joindf2, df1Joindf2("aaa") === df1("a")).select(df1Joindf2("aa"),
+          df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+        val join2 = proj2.child.asInstanceOf[Join]
+        assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
+        assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
+      }
+    })
   }
 
   test("SPARK-47217: deduplication in nested joins without join attribute aliased") {
-    val df1 = Seq((1, 2)).toDF("a", "b")
-    val df2 = Seq((1, 2)).toDF("aa", "bb")
-    val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"), df2("aa"), df1("b"))
-
-    assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("a") === df1("a")),
-      Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
-
-    assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("a") === df1("a")),
-      Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
-
-    val proj1 = df1Joindf2.join(df1, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
-      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
-    val join1 = proj1.child.asInstanceOf[Join]
-    assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
-    assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+    Seq(true, false).foreach(fail => {
+      withSQLConf(
+        SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED.key -> fail.toString) {
+        val df1 = Seq((1, 2)).toDF("a", "b")
+        val df2 = Seq((1, 2)).toDF("aa", "bb")
+        val df1Joindf2 = df1.join(df2, df1("a") === df2("aa")).select(df1("a"), df2("aa"), df1("b"))
+
+        assertCorrectResolution(df1Joindf2.join(df1, df1Joindf2("a") === df1("a")),
+          Resolution.LeftConditionToLeftLeg, Resolution.RightConditionToRightLeg)
 
-    val proj2 = df1.join(df1Joindf2, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
-      df1("a")).queryExecution.analyzed.asInstanceOf[Project]
-    val join2 = proj2.child.asInstanceOf[Join]
-    assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
-    assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
+        assertCorrectResolution(df1.join(df1Joindf2, df1Joindf2("a") === df1("a")),
+          Resolution.LeftConditionToRightLeg, Resolution.RightConditionToLeftLeg)
+
+        val proj1 = df1Joindf2.join(df1, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
+          df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+        val join1 = proj1.child.asInstanceOf[Join]
+        assert(proj1.projectList(0).references.subsetOf(join1.left.outputSet))
+        assert(proj1.projectList(1).references.subsetOf(join1.right.outputSet))
+
+        val proj2 = df1.join(df1Joindf2, df1Joindf2("a") === df1("a")).select(df1Joindf2("a"),
+          df1("a")).queryExecution.analyzed.asInstanceOf[Project]
+        val join2 = proj2.child.asInstanceOf[Join]
+        assert(proj2.projectList(0).references.subsetOf(join2.right.outputSet))
+        assert(proj2.projectList(1).references.subsetOf(join2.left.outputSet))
+      }
+    })
   }
 }
 

From 77b3201f444ac044f16bb19b8e6e4e0266dfe00c Mon Sep 17 00:00:00 2001
From: ashahid <ashahid@cloudera.com>
Date: Fri, 29 Mar 2024 15:41:36 -0700
Subject: [PATCH 22/22] SPARK-47320. cleaned up the code, made the addition of
 DataSet_ID_Tag to LogicalPlan irrespective of boolean
 FAIL_AMBIGUOUS_SELF_JOIN_ENABLED enabled or not.

---
 .../analysis/ColumnResolutionHelper.scala      |  2 +-
 .../catalyst/plans/logical/LogicalPlan.scala   |  2 +-
 .../scala/org/apache/spark/sql/Dataset.scala   | 18 ++++++------------
 3 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index abc59a69bfed9..7bf387cf2bf38 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -537,7 +537,7 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       var currentLp = lp
       var depth = 0
       while (true) {
-        if (currentLp.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG).exists(
+        if (currentLp.getTagValue(LogicalPlan.DATASET_ID_TAG).exists(
           _.contains(datasetId))) {
           return Option(currentLp, depth)
         } else {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
index 1ac3fda25979b..ad3c84e672a48 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -200,7 +200,7 @@ object LogicalPlan {
   //    to the old code path.
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
-  private[spark] val DATASET_RESOLUTION_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
+  private[spark] val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
   private[spark] val UNRESOLVED_ATTRIBUTE_MD_TAG = TreeNodeTag[AttributeReference]("orig-attr")
   private[spark] val DATASET_ID_KEY = "__dataset_id"
   private[spark] val COL_POS_KEY = "__col_position"
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 4b7086432be2f..b1f58baae13b4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -20,7 +20,6 @@ package org.apache.spark.sql
 import java.io.{ByteArrayOutputStream, CharArrayWriter, DataOutputStream}
 
 import scala.annotation.varargs
-import scala.collection.mutable
 import scala.collection.mutable.{ArrayBuffer, HashSet}
 import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
@@ -39,7 +38,7 @@ import org.apache.spark.api.r.RRDD
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.resource.ResourceProfile
-import org.apache.spark.sql.Dataset.DATASET_ID_KEY
+import org.apache.spark.sql.Dataset.{DATASET_ID_KEY, DATASET_ID_TAG}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, QueryPlanningTracker, ScalaReflection, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
@@ -49,7 +48,7 @@ import org.apache.spark.sql.catalyst.json.{JacksonGenerator, JSONOptions}
 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.trees.{TreeNodeTag, TreePattern}
+import org.apache.spark.sql.catalyst.trees.TreePattern
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, IntervalUtils}
 import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId
@@ -75,7 +74,7 @@ private[sql] object Dataset {
   val curId = new java.util.concurrent.atomic.AtomicLong()
   val DATASET_ID_KEY = LogicalPlan.DATASET_ID_KEY
   val COL_POS_KEY = LogicalPlan.COL_POS_KEY
-  val DATASET_ID_TAG = TreeNodeTag[mutable.HashSet[Long]]("dataset_id")
+  val DATASET_ID_TAG = LogicalPlan.DATASET_ID_TAG
 
   def apply[T: Encoder](sparkSession: SparkSession, logicalPlan: LogicalPlan): Dataset[T] = {
     val dataset = new Dataset(sparkSession, logicalPlan, implicitly[Encoder[T]])
@@ -224,14 +223,9 @@ class Dataset[T] private[sql](
 
   @transient private[sql] val logicalPlan: LogicalPlan = {
     val plan = queryExecution.commandExecuted
-    if (sparkSession.conf.get(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED)) {
-      val dsIds = plan.getTagValue(Dataset.DATASET_ID_TAG).getOrElse(new HashSet[Long])
-      dsIds.add(id)
-      plan.setTagValue(Dataset.DATASET_ID_TAG, dsIds)
-    }
     val dsIds = plan.getTagValue(Dataset.DATASET_ID_TAG).getOrElse(new HashSet[Long])
     dsIds.add(id)
-    plan.setTagValue(LogicalPlan.DATASET_RESOLUTION_TAG, dsIds)
+    plan.setTagValue(Dataset.DATASET_ID_TAG, dsIds)
     plan
   }
 
@@ -1181,8 +1175,8 @@ class Dataset[T] private[sql](
       Join(logicalPlan, right.logicalPlan,
         JoinType(joinType), None, JoinHint.NONE)).queryExecution.analyzed.asInstanceOf[Join]
 
-    val leftTagIdMap = planPart1.left.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG)
-    val rightTagIdMap = planPart1.right.getTagValue(LogicalPlan.DATASET_RESOLUTION_TAG)
+    val leftTagIdMap = planPart1.left.getTagValue(DATASET_ID_TAG)
+    val rightTagIdMap = planPart1.right.getTagValue(DATASET_ID_TAG)
 
     val joinExprsRectified = joinExprs.map(_.expr transformUp {
       case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) =>