Address code review: extend hasSameField to compare all fields and simplify null direction handling

Signed-off-by: Selina Song <[email protected]>

Author: Selina Song

core/src/main/java/org/opensearch/sql/calcite/rule/SortReverseOptimizationRule.java

@@ -70,30 +70,48 @@ public void onMatch(RelOptRuleCall call) {
   }
 
   private boolean hasSameFieldWithOppositeDirection(LogicalSort outerSort, LogicalSort innerSort) {
-    if (outerSort.getCollation().getFieldCollations().isEmpty()
-        || innerSort.getCollation().getFieldCollations().isEmpty()) {
+    var outerFields = outerSort.getCollation().getFieldCollations();
+    var innerFields = innerSort.getCollation().getFieldCollations();
+
+    if (outerFields.isEmpty() || innerFields.isEmpty()) {
       LOG.debug("No field collations found");
       return false;
     }
 
-    var outerField = outerSort.getCollation().getFieldCollations().get(0);
-    var innerField = innerSort.getCollation().getFieldCollations().get(0);
-
-    LOG.debug(
-        "Outer field: index={}, direction={}",
-        outerField.getFieldIndex(),
-        outerField.getDirection());
-    LOG.debug(
-        "Inner field: index={}, direction={}",
-        innerField.getFieldIndex(),
-        innerField.getDirection());
-
-    boolean sameField = outerField.getFieldIndex() == innerField.getFieldIndex();
-    boolean oppositeDirection = outerField.getDirection() != innerField.getDirection();
+    // Must have same number of fields
+    if (outerFields.size() != innerFields.size()) {
+      LOG.debug(
+          "Different number of sort fields: outer={}, inner={}",
+          outerFields.size(),
+          innerFields.size());
+      return false;
+    }
 
-    LOG.debug("Same field: {}, Opposite direction: {}", sameField, oppositeDirection);
+    // Check all fields have same index but opposite directions
+    for (int i = 0; i < outerFields.size(); i++) {
+      var outerField = outerFields.get(i);
+      var innerField = innerFields.get(i);
+
+      LOG.debug(
+          "Field {}: outer(index={}, direction={}), inner(index={}, direction={})",
+          i,
+          outerField.getFieldIndex(),
+          outerField.getDirection(),
+          innerField.getFieldIndex(),
+          innerField.getDirection());
+
+      if (outerField.getFieldIndex() != innerField.getFieldIndex()
+          || outerField.getDirection() == innerField.getDirection()) {
+        LOG.debug(
+            "Field {} mismatch: same index={}, opposite direction={}",
+            i,
+            outerField.getFieldIndex() == innerField.getFieldIndex(),
+            outerField.getDirection() != innerField.getDirection());
+        return false;
+      }
+    }
 
-    // Must be same field with opposite directions
-    return sameField && oppositeDirection;
+    LOG.debug("All fields match with opposite directions");
+    return true;
   }
 }

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -411,13 +411,12 @@ public static RelCollation reverseCollation(RelCollation original) {
       RelFieldCollation.Direction reversedDirection = field.direction.reverse();
 
       // Handle null direction properly - reverse it as well
-      RelFieldCollation.NullDirection reversedNullDirection = field.nullDirection;
-      if (reversedNullDirection == RelFieldCollation.NullDirection.FIRST) {
-        reversedNullDirection = RelFieldCollation.NullDirection.LAST;
-      } else if (reversedNullDirection == RelFieldCollation.NullDirection.LAST) {
-        reversedNullDirection = RelFieldCollation.NullDirection.FIRST;
-      }
-      // UNSPECIFIED remains UNSPECIFIED
+      RelFieldCollation.NullDirection reversedNullDirection =
+          field.nullDirection == RelFieldCollation.NullDirection.FIRST
+              ? RelFieldCollation.NullDirection.LAST
+              : field.nullDirection == RelFieldCollation.NullDirection.LAST
+                  ? RelFieldCollation.NullDirection.FIRST
+                  : field.nullDirection;
 
       RelFieldCollation reversedField =
           new RelFieldCollation(field.getFieldIndex(), reversedDirection, reversedNullDirection);