From 2de759364655b263c176dffc4966a50c95a87440 Mon Sep 17 00:00:00 2001
From: huaxingao <huaxin.gao11@gmail.com>
Date: Mon, 23 Sep 2024 10:06:59 -0700
Subject: [PATCH] address comments

---
 .../apache/comet/serde/QueryPlanSerde.scala   | 38 ++++++++-----------
 1 file changed, 15 insertions(+), 23 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 460c80924..4744c2a1f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -3303,33 +3303,25 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     if (partitionSpec.length != orderSpec.length) {
       withInfo(op, "Partitioning and sorting specifications do not match")
       return false
-    } else {
-      val partitionColumnNames = partitionSpec.collect { case a: AttributeReference =>
-        a.name
-      }
-
-      if (partitionColumnNames.length != partitionSpec.length) {
-        withInfo(op, "Unsupported partitioning specification")
-        return false
-      }
-
-      val orderColumnNames = orderSpec.collect { case s: SortOrder =>
-        s.child match {
-          case a: AttributeReference => a.name
-        }
-      }
+    }
 
-      if (orderColumnNames.length != orderSpec.length) {
-        withInfo(op, "Unsupported SortOrder")
-        return false
-      }
+    val partitionColumnNames = partitionSpec.collect { case a: AttributeReference =>
+      a.name
+    }
 
-      if (partitionColumnNames.toSet != orderColumnNames.toSet) {
-        withInfo(op, "Partitioning and sorting specifications do not match")
-        return false
+    val orderColumnNames = orderSpec.collect { case s: SortOrder =>
+      s.child match {
+        case a: AttributeReference => a.name
       }
+    }
 
-      true
+    if (partitionColumnNames.zip(orderColumnNames).exists { case (partCol, orderCol) =>
+        partCol != orderCol
+      }) {
+      withInfo(op, "Partitioning and sorting specifications must be the same.")
+      return false
     }
+
+    true
   }
 }