[SPARK-56917][TESTS][CONNECT] Expand Connect-specific tests for DataFrame column resolution

python/pyspark/sql/tests/connect/test_parity_column.py

@@ -17,6 +17,8 @@
 
 import unittest
 
+from pyspark.errors import AnalysisException
+from pyspark.sql import functions as sf
 from pyspark.sql.tests.test_column import ColumnTestsMixin
 from pyspark.testing.connectutils import ReusedConnectTestCase
 
@@ -38,6 +40,16 @@ def tearDownClass(cls):
     def test_validate_column_types(self):
         super().test_validate_column_types()
 
+    def test_resolve_after_union(self):
+        # Connect diverges from Classic here: Union is treated as a leaf when
+        # walking the plan tree for plan-id resolution, so the left-side plan
+        # id is never found and CANNOT_RESOLVE_DATAFRAME_COLUMN is thrown
+        # before any name-based fallback - in both strict and lenient modes.
+        df1 = self.spark.sql("SELECT 1 AS c")
+        df2 = self.spark.sql("SELECT 2 AS c")
+        with self.assertRaisesRegex(AnalysisException, "CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+            df1.union(df2).select(df1.c).collect()
+
     def test_df_col_resolution_mode(self):
         self.assertEqual(
             self.spark.conf.get("spark.sql.analyzer.strictDataFrameColumnResolution"),
@@ -68,6 +80,30 @@ def test_df_col_resolution_mode(self):
             "false",
         )
 
+    # The shadowing trio diverges in lenient mode: where Classic and Connect
+    # strict raise, lenient resolves the tagged reference by name against the
+    # current (shadowed) output.
+
+    def test_resolve_after_chained_withcolumn_shadow(self):
+        df = self.spark.sql("SELECT 1 AS c")
+        rows = (
+            df.withColumn("c", sf.col("c").cast("string"))
+            .withColumn("c", sf.col("c").cast("int"))
+            .select(df.c)
+            .collect()
+        )
+        self.assertEqual([r.c for r in rows], [1])
+
+    def test_resolve_after_select_alias_shadow(self):
+        df = self.spark.sql("SELECT 1 AS c")
+        rows = df.select(df.c.cast("string").alias("c")).select(df.c).collect()
+        self.assertEqual([r.c for r in rows], ["1"])
+
+    def test_resolve_after_agg_alias_shadow(self):
+        df = self.spark.sql("SELECT 1 AS c")
+        rows = df.groupBy().agg(sf.sum("c").alias("c")).select(df.c).collect()
+        self.assertEqual([r.c for r in rows], [1])
+
 
 if __name__ == "__main__":
     from pyspark.testing import main