Fix IllegalStateException in UnifiedQueryPlanner.preserveCollation on composite-collation plans

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -7,9 +7,10 @@
 
 import static org.opensearch.sql.monitor.profile.MetricName.ANALYZE;
 
+import java.util.List;
 import lombok.RequiredArgsConstructor;
 import org.apache.calcite.rel.RelCollation;
-import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelCollationTraitDef;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelRoot;
 import org.apache.calcite.rel.core.Sort;
@@ -154,9 +155,15 @@ public RelNode plan(String query) {
     }
 
     private RelNode preserveCollation(RelNode logical) {
-      RelCollation collation = logical.getTraitSet().getCollation();
-      if (!(logical instanceof Sort) && collation != RelCollations.EMPTY) {
-        return LogicalSort.create(logical, collation, null, null);
+      if (logical instanceof Sort) {
+        return logical;
+      }
+      List<RelCollation> collations =
+          logical.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE);
+      if (collations != null
+          && collations.size() == 1
+          && !collations.get(0).getFieldCollations().isEmpty()) {
+        return LogicalSort.create(logical, collations.get(0), null, null);
       }
       return logical;
     }

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java

@@ -5,11 +5,14 @@
 
 package org.opensearch.sql.api;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import java.util.Map;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.runtime.CalciteException;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.impl.AbstractSchema;
@@ -198,4 +201,38 @@ public void testPPLPatternsPicksUpDefaults() {
         .assertPlanContains("REGEXP_REPLACE")
         .assertFields("id", "name", "age", "department", "patterns_field");
   }
+
+  @Test
+  public void preserveCollationReturnsExistingSortUnchanged() {
+    var assertion = givenQuery("source = catalog.employees | sort age");
+    assertTrue("top node should be the user's Sort", assertion.plan() instanceof Sort);
+    assertion.assertPlan(
+        "LogicalSort(sort0=[$2], dir0=[ASC-nulls-first])\n"
+            + "  LogicalTableScan(table=[[catalog, employees]])\n");
+  }
+
+  @Test
+  public void preserveCollationLeavesCompositeCollationUnwrapped() {
+    var assertion = givenQuery("source = catalog.employees | sort age | eval x = age");
+    assertFalse(
+        "composite-collation plan must not be re-wrapped in a Sort",
+        assertion.plan() instanceof Sort);
+    assertion.assertFields("id", "name", "age", "department", "x");
+  }
+
+  @Test
+  public void preserveCollationMaterializesSingleDerivedCollation() {
+    givenQuery("makeresults format=csv data='name\nJohn\nSarah'")
+        .assertPlan(
+            "LogicalSort(sort0=[$0], dir0=[ASC])\n"
+                + "  LogicalProject(name=[CAST($0):VARCHAR NOT NULL])\n"
+                + "    LogicalValues(tuples=[[{ 'John' }, { 'Sarah' }]])\n");
+  }
+
+  @Test
+  public void preserveCollationLeavesUnsortedPlanUnwrapped() {
+    var assertion = givenQuery("source = catalog.employees");
+    assertFalse("unsorted plan must not be wrapped in a Sort", assertion.plan() instanceof Sort);
+    assertion.assertPlan("LogicalTableScan(table=[[catalog, employees]])\n");
+  }
 }

api/src/test/java/org/opensearch/sql/api/compiler/UnifiedQueryCompilerTest.java

@@ -57,6 +57,27 @@ public void testComplexQuery() throws Exception {
     }
   }
 
+  @Test
+  public void testCompositeCollationQueryExecutes() throws Exception {
+    RelNode plan = planner.plan("source = catalog.employees | sort age | eval x = age");
+    try (PreparedStatement statement = compiler.compile(plan)) {
+      ResultSet resultSet = statement.executeQuery();
+
+      verify(resultSet)
+          .expectSchema(
+              col("id", INTEGER),
+              col("name", VARCHAR),
+              col("age", INTEGER),
+              col("department", VARCHAR),
+              col("x", INTEGER))
+          .expectData(
+              row(1, "Alice", 25, "Engineering", 25),
+              row(2, "Bob", 35, "Sales", 35),
+              row(3, "Charlie", 45, "Engineering", 45),
+              row(4, "Diana", 28, "Marketing", 28));
+    }
+  }
+
   @Test(expected = IllegalStateException.class)
   public void testCompileFailure() {
     RelNode mockPlan = Mockito.mock(RelNode.class);

integ-test/src/test/java/org/opensearch/sql/api/UnifiedQueryOpenSearchIT.java

@@ -104,6 +104,30 @@ public void testMultiplePPLQueryExecutionWithSameContext() throws Exception {
     }
   }
 
+  @Test
+  public void testSortThenEvalCopyExecutesEndToEnd() throws Exception {
+    String pplQuery =
+        String.format(
+            "source = opensearch.%s | sort age | eval age_copy = age"
+                + " | fields firstname, age, age_copy",
+            TEST_INDEX_ACCOUNT);
+
+    RelNode plan = planner.plan(pplQuery);
+    try (PreparedStatement statement = compiler.compile(plan)) {
+      ResultSet resultSet = statement.executeQuery();
+
+      verify(resultSet)
+          .expectSchema(col("firstname", VARCHAR), col("age", BIGINT), col("age_copy", BIGINT));
+
+      int rows = 0;
+      while (resultSet.next()) {
+        rows++;
+        assertEquals("age_copy must mirror age", resultSet.getObject(2), resultSet.getObject(3));
+      }
+      assertTrue("expected at least one row", rows > 0);
+    }
+  }
+
   /**
    * Creates a dynamic schema that creates OpenSearchIndex on-demand for any table name. This allows
    * querying any index without pre-registering it.