[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

@@ -49,7 +49,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._`.
@@ -1701,6 +1701,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)) =>
@@ -2085,6 +2086,29 @@ 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).
+     */
+    private def retainExceptedColumnsAsHiddenOutput(original: Project, expanded: Project): Unit = {
+      val retain = original.projectList.exists {
+        case s: UnresolvedStarExceptOrReplace => s.retainExceptedColumnsAsHidden
+        case _ => false
+      }
+      if (retain) {
+        val excepted = expanded.child.output.filterNot(expanded.outputSet.contains)
+        if (excepted.nonEmpty) {
+          expanded.setTagValue(
+            Project.hiddenOutputTag,
+            excepted.map(_.markAsQualifiedAccessOnly()) ++ expanded.child.metadataOutput)
+        }
+      }
+    }
+
     /**
      * Returns true if `exprs` contains a [[Star]].
      */

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

@@ -640,11 +640,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

@@ -8158,11 +8158,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/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out

@@ -992,6 +992,65 @@ Project [a#x, a#x, z2#x]
                                              +- LocalRelation [a#x]
 
 
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select t.a
+-- !query analysis
+Project [a#x]
++- Project [(a#x + 1) AS a#x, b#x, a#x]
+   +- SubqueryAlias t
+      +- LocalRelation [a#x, b#x]
+
+
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select a, t.a, t.b
+-- !query analysis
+Project [a#x, a#x, b#x]
++- Project [(a#x + 1) AS a#x, b#x, a#x]
+   +- SubqueryAlias t
+      +- LocalRelation [a#x, b#x]
+
+
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select *
+-- !query analysis
+Project [a#x, b#x]
++- Project [(a#x + 1) AS a#x, b#x]
+   +- SubqueryAlias t
+      +- LocalRelation [a#x, b#x]
+
+
+-- !query
+table t
+|> as u
+|> set x = x + 1
+|> select x, u.x
+-- !query analysis
+Project [x#x, x#x]
++- Project [(x#x + 1) AS x#x, y#x, x#x]
+   +- SubqueryAlias u
+      +- SubqueryAlias spark_catalog.default.t
+         +- Relation spark_catalog.default.t[x#x,y#x] csv
+
+
+-- !query
+values (1) as t(a)
+|> set a = a + 1
+|> set a = a + 1
+|> select a, t.a
+-- !query analysis
+Project [a#x, a#x]
++- Project [(a#x + 1) AS a#x, a#x]
+   +- Project [(a#x + 1) AS a#x, a#x]
+      +- SubqueryAlias t
+         +- LocalRelation [a#x]
+
+
 -- !query
 table t
 |> set z = 1

sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql

@@ -361,6 +361,33 @@ values (0), (1) lhs(a)
 |> limit 2
 |> select lhs.a, rhs.a, z2;
 
+-- A table alias still refers to the original row value of a column affected by SET.
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select t.a;
+
+-- The unqualified name refers to the assigned value and the qualified name to the original one.
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select a, t.a, t.b;
+
+-- The retained source column stays out of the output schema unless it is named explicitly.
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select *;
+
+-- Qualified access to a column affected by SET, through an alias added by the AS operator.
+table t
+|> as u
+|> set x = x + 1
+|> select x, u.x;
+
+-- The alias keeps pointing at the original source value across a sequence of SET operators.
+values (1) as t(a)
+|> set a = a + 1
+|> set a = a + 1
+|> select a, t.a;
+
 -- SET operators: negative tests.
 ---------------------------------
 

sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out

@@ -909,6 +909,59 @@ struct<a:int,a:int,z2:int>
 1	1	4
 
 
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select t.a
+-- !query schema
+struct<a:int>
+-- !query output
+1
+
+
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select a, t.a, t.b
+-- !query schema
+struct<a:int,a:int,b:int>
+-- !query output
+2	1	10
+
+
+-- !query
+values (1, 10) as t(a, b)
+|> set a = a + 1
+|> select *
+-- !query schema
+struct<a:int,b:int>
+-- !query output
+2	10
+
+
+-- !query
+table t
+|> as u
+|> set x = x + 1
+|> select x, u.x
+-- !query schema
+struct<x:int,x:int>
+-- !query output
+1	0
+2	1
+
+
+-- !query
+values (1) as t(a)
+|> set a = a + 1
+|> set a = a + 1
+|> select a, t.a
+-- !query schema
+struct<a:int,a:int>
+-- !query output
+3	1
+
+
 -- !query
 table t
 |> set z = 1