[SPARK-59146][SQL] Retain qualified access to source columns affected by pipe SET

docs/sql-pipe-syntax.md

@@ -289,14 +289,15 @@ VALUES (0), (1) tab(col)
 +---+
 
 VALUES (0), (1) tab(col)
-|> SET col = col * 2;
+|> SET col = col * 2
+|> SELECT col, tab.col;
 
-+---+
-|col|
-+---+
-|  0|
-|  2|
-+---+
++---+---+
+|col|col|
++---+---+
+|  0|  0|
+|  2|  1|
++---+---+
 ```
 
 #### DROP

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

@@ -48,7 +48,7 @@ import org.apache.spark.sql.catalyst.trees.AlwaysProcess
 import org.apache.spark.sql.catalyst.trees.CurrentOrigin.withOrigin
 import org.apache.spark.sql.catalyst.trees.TreePattern._
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
-import org.apache.spark.sql.catalyst.util.{toPrettySQL, trimTempResolvedColumn, CharVarcharUtils, GeneratedColumn}
+import org.apache.spark.sql.catalyst.util.{toPrettySQL, trimTempResolvedColumn, CharVarcharUtils, GeneratedColumn, MetadataColumnHelper}
 import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns._
 // `View` is aliased to `V2View` to avoid clashing with the logical-plan `View` imported via
 // `org.apache.spark.sql.catalyst.plans.logical._`.
@@ -1163,28 +1163,34 @@ class Analyzer(
     }
 
     private def getMetadataAttributes(plan: LogicalPlan): Seq[Attribute] = {
+      val childrenMetadataOutput = childrenMetadataOutputById(plan)
       plan.expressions.flatMap(_.collect {
         case a: Attribute if a.isMetadataCol => a
-        case a: Attribute
-          if plan.children.exists(c => c.metadataOutput.exists(_.exprId == a.exprId)) =>
-          plan.children.collectFirst {
-            case c if c.metadataOutput.exists(_.exprId == a.exprId) =>
-              c.metadataOutput.find(_.exprId == a.exprId).get
-          }.get
+        case a: Attribute if childrenMetadataOutput.contains(a.exprId) =>
+          childrenMetadataOutput(a.exprId)
       })
     }
 
     private def hasMetadataCol(plan: LogicalPlan): Boolean = {
+      lazy val childrenMetadataOutput = childrenMetadataOutputById(plan)
       plan.expressions.exists(_.exists {
         case a: Attribute =>
           // If an attribute is resolved before being labeled as metadata
           // (i.e. from the originating Dataset), we check with expression ID
-          a.isMetadataCol ||
-            plan.children.exists(c => c.metadataOutput.exists(_.exprId == a.exprId))
+          a.isMetadataCol || childrenMetadataOutput.contains(a.exprId)
         case _ => false
       })
     }
 
+    // The metadata output of the children, keeping the first attribute for each expression ID.
+    // Looking attributes up by ID keeps this rule linear when a child's metadata output is as wide
+    // as its output, as the hidden source row of the SQL pipe SET operator is.
+    private def childrenMetadataOutputById(plan: LogicalPlan): collection.Map[ExprId, Attribute] = {
+      val byId = mutable.HashMap.empty[ExprId, Attribute]
+      plan.children.foreach(_.metadataOutput.foreach(a => byId.getOrElseUpdate(a.exprId, a)))
+      byId
+    }
+
     private def addMetadataCol(
         plan: LogicalPlan,
         requiredAttrIds: Set[ExprId],
@@ -1783,6 +1789,7 @@ class Analyzer(
         if (expanded.projectList.size < p.projectList.size) {
           checkTrailingCommaInSelect(expanded, starRemoved = true)
         }
+        retainExceptedColumnsAsHiddenOutput(p, expanded)
         expanded
       // If the filter list contains Stars, expand it.
       case p: Filter if containsStar(Seq(p.condition)) =>
@@ -2182,6 +2189,50 @@ class Analyzer(
       }.map(_.asInstanceOf[NamedExpression])
     }
 
+    /**
+     * The SQL pipe SET operator is implemented as a star expansion that excludes the assigned
+     * column and appends a replacement of the same name. That drops the original attribute from
+     * the project list, which would also make it unreachable through its table alias, contradicting
+     * the documented behavior that table aliases keep referring to the original row values after an
+     * assignment. Retain the excluded attributes as hidden output instead, the same way USING joins
+     * hide their duplicated join keys (SPARK-59146).
+     *
+     * The hidden output holds the whole qualified source row, with every position in its original
+     * order, not just the excluded attributes, because a qualified star expands hidden output
+     * ahead of the visible output. Retaining only the excluded attributes would move them to the
+     * front, so that `VALUES (1, 2, 3) AS t(a, b, c) |> SET b = 20 |> SELECT t.*` would return
+     * `(b, a, c)`. Attributes that are also visible are emitted once by the star expansion, in
+     * their hidden position. The Project does not hide them, so they are also recorded in
+     * [[Project.visibleHiddenOutputTag]], which keeps them regular in DataFrame column resolution.
+     */
+    private def retainExceptedColumnsAsHiddenOutput(original: Project, expanded: Project): Unit = {
+      val retain = original.projectList.exists {
+        case s: UnresolvedStarExceptOrReplace => s.retainExceptedColumnsAsHidden
+        case _ => false
+      }
+      if (retain) {
+        val child = expanded.child
+        if (child.output.exists(!expanded.outputSet.contains(_))) {
+          // The row a qualified star sees on the child, keeping every one of its positions. Only
+          // qualified attributes are reachable through a qualified star.
+          val sourceRow = UnresolvedStarBase
+            .mergeHiddenAndVisibleOutput(
+              child.metadataOutput.filter(_.qualifiedAccessOnly), child.output)
+            .filter(_.qualifier.nonEmpty)
+          val sourceRowIds = sourceRow.map(_.exprId).toSet
+          // The rule forwards the original tags, such as the Spark Connect plan id, only onto a
+          // node without tags, so copy them here before adding the hidden output tag.
+          expanded.copyTagsFrom(original)
+          expanded.setTagValue(
+            Project.hiddenOutputTag,
+            sourceRow.map(_.markAsQualifiedAccessOnly()) ++
+              child.metadataOutput.filterNot(a => sourceRowIds.contains(a.exprId)))
+          expanded.setTagValue(
+            Project.visibleHiddenOutputTag, sourceRow.filter(expanded.outputSet.contains))
+        }
+      }
+    }
+
     /**
      * Returns true if `exprs` contains a [[Star]].
      */

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

@@ -701,8 +701,12 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
     //   - `AddMetadataColumns` adds a hidden key to the `output` of the USING join
     //     wrapper that hides it, and a metadata column to the `output` of a relation.
     //     That node introduces the attribute in its `metadataOutput`, so it stays hidden.
+    //   - The Project of the SQL pipe SET operator introduces the whole source row in its
+    //     `metadataOutput`, to keep the column order of a qualified star, but only hides the
+    //     columns it assigns. The others are in `Project.visibleHiddenOutputTag` and visible.
     val introducedHidden = AttributeSet(p.metadataOutput) --
-      p.children.flatMap(_.metadataOutput)
+      p.children.flatMap(_.metadataOutput) --
+      p.getTagValue(Project.visibleHiddenOutputTag).getOrElse(Nil)
     val visible = p.outputSet -- introducedHidden
     val filtered = candidates.flatMap { c =>
       val hidden = c.hidden || !c.expr.references.subsetOf(visible)

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
+import scala.collection.mutable
+
 import org.apache.spark.SparkException
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow, TableIdentifier}
@@ -593,7 +595,8 @@ trait UnresolvedStarBase extends Star with Unevaluable {
       // keep any restrictions that may break column resolution for normal attributes.
       // See SPARK-42084 for more details.
       .map(_.markAsAllowAnyAccess())
-    val expandedAttributes = (hiddenOutput ++ parameters.childOperatorOutput)
+    val expandedAttributes = UnresolvedStarBase
+      .mergeHiddenAndVisibleOutput(hiddenOutput, parameters.childOperatorOutput)
       .filter(matchedQualifier(_, target.get, parameters.resolver))
 
     if (expandedAttributes.nonEmpty) return expandedAttributes
@@ -635,6 +638,51 @@ trait UnresolvedStarBase extends Star with Unevaluable {
   override def toString: String = target.map(_.mkString("", ".", ".")).getOrElse("") + "*"
 }
 
+object UnresolvedStarBase {
+
+  /**
+   * Merges the hidden output of a plan with its visible output, in the order a qualified star
+   * expands them.
+   *
+   * Hidden output usually holds attributes that are not visible at all, such as the duplicated
+   * join keys of a USING join. It may also repeat an attribute that is visible, to pin that
+   * attribute's position in the expansion, as the SQL pipe SET operator does to keep the original
+   * column order of the row its table alias refers to (SPARK-59146).
+   *
+   * Every occurrence of the source row is therefore preserved: a hidden attribute consumes at
+   * most one visible occurrence, which is emitted in the hidden attribute's place, because a
+   * projection may repeat the same attribute. Visible attributes that no hidden attribute
+   * consumed keep their relative order, after the hidden output.
+   *
+   * Both outputs can be as wide as the source row, for example under a wide USING join or a pipe
+   * SET, so the merge runs in time linear in their combined size.
+   */
+  def mergeHiddenAndVisibleOutput(
+      hiddenOutput: Seq[Attribute],
+      visibleOutput: Seq[Attribute]): Seq[Attribute] = {
+    if (hiddenOutput.isEmpty) return visibleOutput
+    val visible = visibleOutput.toIndexedSeq
+    // The positions of the occurrences of each visible attribute, in order. Built back to front so
+    // that every list is in ascending order and its head is the first unconsumed occurrence.
+    val occurrences = mutable.HashMap.empty[ExprId, List[Int]]
+    visible.indices.reverseIterator.foreach { i =>
+      val exprId = visible(i).exprId
+      occurrences(exprId) = i :: occurrences.getOrElse(exprId, Nil)
+    }
+    val consumed = new Array[Boolean](visible.length)
+    val merged = hiddenOutput.map { hidden =>
+      occurrences.get(hidden.exprId) match {
+        case Some(i :: remaining) =>
+          occurrences(hidden.exprId) = remaining
+          consumed(i) = true
+          visible(i)
+        case _ => hidden
+      }
+    }
+    merged ++ visible.indices.collect { case i if !consumed(i) => visible(i) }
+  }
+}
+
 /**
  * Represents some of the input attributes to a given relational operator, for example in
  * "SELECT * EXCEPT(a) FROM ...".
@@ -649,11 +697,20 @@ trait UnresolvedStarBase extends Star with Unevaluable {
  *                     expressions removed by EXCEPT. If present, the length of this list must
  *                     be the same as the length of the EXCEPT list. This supports replacing
  *                     expressions instead of excluding them from the original SELECT list.
+ *
+ * @param retainExceptedColumnsAsHidden if true, the excluded attributes are kept as the hidden
+ *                                      output of the enclosing [[Project]], so that they remain
+ *                                      reachable through their table alias. The SQL pipe SET
+ *                                      operator sets this, since it documents that table aliases
+ *                                      keep referring to the original row values after an
+ *                                      assignment. It stays false for SELECT * EXCEPT, where the
+ *                                      excluded columns must not be reachable at all.
  */
 case class UnresolvedStarExceptOrReplace(
     target: Option[Seq[String]],
     excepts: Seq[Seq[String]],
-    replacements: Option[Seq[NamedExpression]])
+    replacements: Option[Seq[NamedExpression]],
+    retainExceptedColumnsAsHidden: Boolean = false)
   extends LeafExpression with UnresolvedStarBase {
 
   final override val nodePatterns: Seq[TreePattern] = Seq(UNRESOLVED_STAR_EXCEPT_OR_REPLACE)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -8330,11 +8330,16 @@ class AstBuilder extends DataTypeAstBuilder
         // Add an UnresolvedStarExceptOrReplace to exclude the SET expression name from the relation
         // and add the new SET expression to the projection list.
         // Use a PipeSelect expression to make sure it does not contain any aggregate functions.
+        // Retain the excluded source column as hidden output so that qualified references such as
+        // `t.a` keep returning the original row value after `|> SET a = ...` (SPARK-59146).
         val replacement =
           Alias(PipeExpression(target, isAggregate = false, PipeOperators.setClause), ident)()
         val projectList: Seq[NamedExpression] =
           Seq(UnresolvedStarExceptOrReplace(
-            target = None, excepts = Seq(Seq(ident)), replacements = Some(Seq(replacement))))
+            target = None,
+            excepts = Seq(Seq(ident)),
+            replacements = Some(Seq(replacement)),
+            retainExceptedColumnsAsHidden = true))
         // Add a projection to implement the SET operator using the UnresolvedStarExceptOrReplace
         // expression. We do this once per SET assignment to allow for multiple SET assignments with
         // optional lateral references to previous ones.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -120,6 +120,14 @@ case class Project(projectList: Seq[NamedExpression], child: LogicalPlan)
 object Project {
   val hiddenOutputTag: TreeNodeTag[Seq[Attribute]] = TreeNodeTag[Seq[Attribute]]("hidden_output")
 
+  /**
+   * The attributes of [[hiddenOutputTag]] that the Project also keeps visible in its output. They
+   * are in the hidden output only to keep their positions in a qualified star, as are the source
+   * columns that the SQL pipe SET operator does not assign, so the Project does not hide them.
+   */
+  val visibleHiddenOutputTag: TreeNodeTag[Seq[Attribute]] =
+    TreeNodeTag[Seq[Attribute]]("visible_hidden_output")
+
   def matchSchema(plan: LogicalPlan, schema: StructType, conf: SQLConf): Project = {
     assert(plan.resolved)
     val projectList =

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala

@@ -2002,6 +2002,114 @@ class AnalysisSuite extends AnalysisTest with Matchers {
     val expectedPlan = Project(Seq(UnresolvedAttribute("i")), addColumnF).analyze
     checkAnalysis(inputPlan, expectedPlan)
   }
+
+  test("SPARK-59146: pipe SET keeps the original Project tags alongside the hidden output") {
+    // The Project that the SQL pipe SET operator builds, as a Spark Connect relation.
+    val set = Project(
+      Seq(UnresolvedStarExceptOrReplace(
+        target = None,
+        excepts = Seq(Seq("a")),
+        replacements = Some(Seq(Alias(Literal(1), "a")())),
+        retainExceptedColumnsAsHidden = true)),
+      testRelation.subquery("t"))
+    set.setTagValue(LogicalPlan.PLAN_ID_TAG, 42L)
+
+    val analyzed = getAnalyzer.execute(set)
+    assert(analyzed.isInstanceOf[Project], analyzed)
+    assert(analyzed.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(42L))
+    assert(analyzed.getTagValue(Project.hiddenOutputTag).exists(_.map(_.name) == Seq("a")))
+  }
+
+  test("SPARK-59146: DataFrame columns that pipe SET did not assign stay visible") {
+    // The Project of `|> SET a = 'x'` over `t(a, b)`, as a Spark Connect relation. Each use of the
+    // DataFrame is planned from a new relation, with new expression IDs and the same plan id.
+    def set(source: LocalRelation): LogicalPlan = {
+      val project = Project(
+        Seq(UnresolvedStarExceptOrReplace(
+          target = None,
+          excepts = Seq(Seq("a")),
+          replacements = Some(Seq(Alias(Literal("x"), "a")())),
+          retainExceptedColumnsAsHidden = true)),
+        source.subquery("t"))
+      project.setTagValue(LogicalPlan.PLAN_ID_TAG, 1L)
+      project
+    }
+    // The DataFrame is used directly, and as the right side of a USING join on `b` that hides its
+    // `b`. SET retains its whole source row as hidden output, but only to keep the column order,
+    // so the direct candidate is regular and wins. Treating it as hidden too would make the
+    // reference ambiguous.
+    val direct = LocalRelation($"a".string, $"b".string)
+    val reused = Join(
+      LocalRelation($"b".string, $"c".int),
+      set(LocalRelation($"a".string, $"b".string)),
+      UsingJoin(Inner, Seq("b")),
+      None,
+      JoinHint.NONE)
+    val column = UnresolvedAttribute("b")
+    column.setTagValue(LogicalPlan.PLAN_ID_TAG, 1L)
+    val plan = Project(Seq(column), Join(set(direct), reused, Inner, None, JoinHint.NONE))
+
+    val output = getAnalyzer.execute(plan).output
+    assert(output.map(_.exprId) == Seq(direct.output(1).exprId), output)
+  }
+
+  test("SPARK-59146: merge hidden and visible output for a qualified star") {
+    // Compares instances, so that the qualifier and metadata of every position are checked too.
+    def checkMerge(
+        hidden: Seq[Attribute],
+        visible: Seq[Attribute],
+        expected: Seq[Attribute]): Unit = {
+      val merged = UnresolvedStarBase.mergeHiddenAndVisibleOutput(hidden, visible)
+      assert(merged.length == expected.length, merged)
+      merged.zip(expected).foreach { case (actual, instance) =>
+        assert(actual eq instance, s"$actual in $merged is not the instance $instance")
+      }
+    }
+    def attributes(qualifier: String, names: String*): Seq[AttributeReference] =
+      names.map(name => AttributeReference(name, IntegerType)(qualifier = Seq(qualifier)))
+
+    // Disjoint, as for the join key that a USING join hides: the hidden output goes first.
+    val Seq(lk, x) = attributes("l", "k", "x")
+    val Seq(rk, y) = attributes("r", "k", "y")
+    val hiddenKey = rk.markAsQualifiedAccessOnly()
+    checkMerge(Seq(hiddenKey), Seq(lk, x, y), Seq(hiddenKey, lk, x, y))
+    checkMerge(Nil, Seq(lk, x, y), Seq(lk, x, y))
+
+    // Overlapping, as for pipe SET: the hidden source row pins the position of every visible
+    // attribute, which is emitted in place of its hidden copy.
+    val Seq(a, b, c) = attributes("t", "a", "b", "c")
+    val assigned = AttributeReference("b", IntegerType)()
+    val row = Seq(a, b, c).map(_.markAsQualifiedAccessOnly())
+    checkMerge(row, Seq(a, c, assigned), Seq(a, row(1), c, assigned))
+
+    // Repeated, as for `SELECT a, a, b`: the occurrences share an ExprId, and are told apart here
+    // by their qualifier. A hidden occurrence consumes the first unconsumed visible one, so every
+    // position is kept exactly once.
+    val repeatedA = a.withQualifier(Seq("u"))
+    val repeatedRow = Seq(a, repeatedA, b).map(_.markAsQualifiedAccessOnly())
+    checkMerge(
+      repeatedRow, Seq(a, repeatedA, assigned), Seq(a, repeatedA, repeatedRow(2), assigned))
+    checkMerge(Seq(row(1), row(0)), Seq(a, repeatedA), Seq(row(1), a, repeatedA))
+    checkMerge(repeatedRow.take(2), Seq(a), Seq(a, repeatedRow(1)))
+  }
+
+  test("SPARK-59146: a qualified star after repeated pipe SETs returns the original row") {
+    // The second SET sees the hidden output of the first one overlap its visible output.
+    def set(column: String, child: LogicalPlan): LogicalPlan = Project(
+      Seq(UnresolvedStarExceptOrReplace(
+        target = None,
+        excepts = Seq(Seq(column)),
+        replacements = Some(Seq(Alias(Literal("x"), column)())),
+        retainExceptedColumnsAsHidden = true)),
+      child)
+    val plan = Project(
+      Seq(UnresolvedStar(Some(Seq("t")))),
+      set("a", set("c", testRelation2.subquery("t"))))
+
+    val output = getAnalyzer.execute(plan).output
+    assert(output.map(_.exprId) == testRelation2.output.map(_.exprId), output)
+    assert(output.forall(_.qualifier == Seq("t")), output)
+  }
 }
 
 /**