address comments

huaxingao · Sep 23, 2024 · 2de7593 · 2de7593
1 parent 0d84fd2
commit 2de7593
Showing 1 changed file with 15 additions and 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
@@ -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
   }
 }