[SPARK-46260][PYTHON][SQL] DataFrame.withColumnsRenamed should respect the dict ordering

python/pyspark/sql/dataframe.py

@@ -6272,7 +6272,18 @@ def withColumnsRenamed(self, colsMap: Dict[str, str]) -> "DataFrame":
                 message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__},
             )
 
-        return DataFrame(self._jdf.withColumnsRenamed(colsMap), self.sparkSession)
+        col_names: List[str] = []
+        new_col_names: List[str] = []
+        for k, v in colsMap.items():
+            col_names.append(k)
+            new_col_names.append(v)
+
+        return DataFrame(
+            self._jdf.withColumnsRenamed(
+                _to_seq(self._sc, col_names), _to_seq(self._sc, new_col_names)
+            ),
+            self.sparkSession,
+        )
 
     def withMetadata(self, columnName: str, metadata: Dict[str, Any]) -> "DataFrame":
         """Returns a new :class:`DataFrame` by updating an existing column with metadata.

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

@@ -77,6 +77,11 @@ def test_to_pandas_from_mixed_dataframe(self):
     def test_toDF_with_string(self):
         super().test_toDF_with_string()
 
+    # TODO(SPARK-46260): DataFrame.withColumnsRenamed should respect the dict ordering
+    @unittest.skip("Fails in Spark Connect, should enable.")
+    def test_ordering_of_with_columns_renamed(self):
+        super().test_ordering_of_with_columns_renamed()
+
 
 if __name__ == "__main__":
     import unittest

python/pyspark/sql/tests/test_dataframe.py

@@ -163,6 +163,15 @@ def test_with_columns_renamed(self):
             message_parameters={"arg_name": "colsMap", "arg_type": "tuple"},
         )
 
+    def test_ordering_of_with_columns_renamed(self):
+        df = self.spark.range(10)
+
+        df1 = df.withColumnsRenamed({"id": "a", "a": "b"})
+        self.assertEqual(df1.columns, ["b"])
+
+        df2 = df.withColumnsRenamed({"a": "b", "id": "a"})
+        self.assertEqual(df2.columns, ["a"])
+
     def test_drop_duplicates(self):
         # SPARK-36034 test that drop duplicates throws a type error when in correct type provided
         df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"])

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -2922,18 +2922,29 @@ class Dataset[T] private[sql](
    */
   @throws[AnalysisException]
   def withColumnsRenamed(colsMap: Map[String, String]): DataFrame = withOrigin {
+    val (colNames, newColNames) = colsMap.toSeq.unzip
+    withColumnsRenamed(colNames, newColNames)
+  }
+
+  private def withColumnsRenamed(
+    colNames: Seq[String],
+    newColNames: Seq[String]): DataFrame = withOrigin {
+    require(colNames.size == newColNames.size,
+      s"The size of existing column names: ${colNames.size} isn't equal to " +
+        s"the size of new column names: ${newColNames.size}")
+
     val resolver = sparkSession.sessionState.analyzer.resolver
     val output: Seq[NamedExpression] = queryExecution.analyzed.output
 
-    val projectList = colsMap.foldLeft(output) {
+    val projectList = colNames.zip(newColNames).foldLeft(output) {
       case (attrs, (existingName, newName)) =>
-      attrs.map(attr =>
-        if (resolver(attr.name, existingName)) {
-          Alias(attr, newName)()
-        } else {
-          attr
-        }
-      )
+        attrs.map(attr =>
+          if (resolver(attr.name, existingName)) {
+            Alias(attr, newName)()
+          } else {
+            attr
+          }
+        )
     }
     SchemaUtils.checkColumnNameDuplication(
       projectList.map(_.name),