Extend to support column name

Author: malthe

python/pyspark/sql/classic/dataframe.py

@@ -1638,8 +1638,10 @@ def withColumns(self, *colsMap: Dict[str, Column]) -> ParentDataFrame:
             self.sparkSession,
         )
 
-    def withColumn(self, colName: str, col: Column) -> ParentDataFrame:
-        if not isinstance(col, Column):
+    def withColumn(self, colName: str, col: "ColumnOrName") -> ParentDataFrame:
+        if isinstance(col, str):
+            col = F.col(col)
+        elif not isinstance(col, Column):
             raise PySparkTypeError(
                 errorClass="NOT_COLUMN",
                 messageParameters={"arg_name": "col", "arg_type": type(col).__name__},

python/pyspark/sql/connect/dataframe.py

@@ -932,8 +932,10 @@ def withColumns(self, *colsMap: Dict[str, Column]) -> ParentDataFrame:
             session=self._session,
         )
 
-    def withColumn(self, colName: str, col: Column) -> ParentDataFrame:
-        if not isinstance(col, Column):
+    def withColumn(self, colName: str, col: "ColumnOrName") -> ParentDataFrame:
+        if isinstance(col, str):
+            col = F.col(col)
+        elif not isinstance(col, Column):
             raise PySparkTypeError(
                 errorClass="NOT_COLUMN",
                 messageParameters={"arg_name": "col", "arg_type": type(col).__name__},

python/pyspark/sql/dataframe.py

@@ -5651,7 +5651,7 @@ def withColumns(self, *colsMap: Dict[str, Column]) -> "DataFrame":
         ...
 
     @dispatch_df_method
-    def withColumn(self, colName: str, col: Column) -> "DataFrame":
+    def withColumn(self, colName: str, col: "ColumnOrName") -> "DataFrame":
         """
         Returns a new :class:`DataFrame` by adding a column or replacing the
         existing column that has the same name.
@@ -5668,8 +5668,8 @@ def withColumn(self, colName: str, col: Column) -> "DataFrame":
         ----------
         colName : str
             string, name of the new column.
-        col : :class:`Column`
-            a :class:`Column` expression for the new column.
+        cols: str or :class:`Column`
+            A name of a column, or a :class:`Column` expression for the new column.
 
         Returns
         -------

python/pyspark/sql/tests/test_dataframe.py

@@ -413,6 +413,10 @@ def test_with_column_with_existing_name(self):
         keys = self.df.withColumn("key", self.df.key).select("key").collect()
         self.assertEqual([r.key for r in keys], list(range(100)))
 
+    def test_with_column_with_column_name(self):
+        keys = self.df.withColumn("key", "key").select("key").collect()
+        self.assertEqual([r.key for r in keys], list(range(100)))
+
     # regression test for SPARK-10417
     def test_column_iterator(self):
         def foo():