Add tests for preserveCollation composite-collation fix; drop comment

Cover UnifiedQueryPlanner#preserveCollation with tests and remove the
now-redundant inline comment (the rationale lives in the commit/PR).

Unit tests (UnifiedQueryPlannerTest) exercise all four branches:
- an existing top-level Sort is returned unchanged (no double-wrap);
- a composite collation is left unwrapped instead of crashing
  (the regression: `sort age | eval x = age` makes x mirror age, so the
  top LogicalProject derives multiple collations stored as a
  RelCompositeTrait; the old getCollation() read threw "Trait index N
  has multiple values", surfacing as a 500);
- a single genuine collation is materialized as a LogicalSort;
- an unsorted plan is returned unwrapped.

Execution coverage: UnifiedQueryCompilerTest compiles and runs the
previously-crashing plan in-memory, and UnifiedQueryOpenSearchIT runs
the same sort+eval-copy shape end to end against a real OpenSearch
index. The composite-collation tests were confirmed to fail against the
pre-fix implementation.

Signed-off-by: Louis Chu <lingzhichu.clz@gmail.com>

Author: noCharger

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

@@ -158,11 +158,6 @@ private RelNode preserveCollation(RelNode logical) {
       if (logical instanceof Sort) {
         return logical;
       }
-      // Only a genuine single-sort collation is materialized. A literal-row plan (makeresults
-      // data=) carries several incidental derived collations kept as a RelCompositeTrait; those
-      // are not a user sort, so getTraits returns more than one and we leave the plan unwrapped.
-      // RelTraitSet.getCollation() would instead cast the composite to a single RelCollation and
-      // throw.
       List<RelCollation> collations =
           logical.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE);
       if (collations != null

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,56 @@ public void testPPLPatternsPicksUpDefaults() {
         .assertPlanContains("REGEXP_REPLACE")
         .assertFields("id", "name", "age", "department", "patterns_field");
   }
+
+  /**
+   * {@code preserveCollation} returns an existing top-level {@link Sort} untouched — it must not
+   * wrap a user sort in a second, redundant {@code LogicalSort}.
+   */
+  @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");
+  }
+
+  /**
+   * Regression test for the composite-collation crash. {@code eval x = age} after {@code sort age}
+   * makes {@code x} mirror {@code age}, so the top {@link
+   * org.apache.calcite.rel.logical.LogicalProject} derives more than one collation, which Calcite
+   * stores as a {@code RelCompositeTrait}. The previous {@code getTraitSet().getCollation()} read
+   * threw ({@code "Trait index 1 has multiple values ... use getTraits"}), surfacing as a 500. A
+   * composite is not a user sort, so the plan must build and be left unwrapped.
+   */
+  @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");
+  }
+
+  /**
+   * A single genuine (non-empty) collation on a non-{@link Sort} top is materialized as a {@code
+   * LogicalSort} so the ordering survives for downstream consumers. Single-column {@code
+   * makeresults} derives exactly one collation; this branch is unchanged by the composite fix.
+   */
+  @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");
+  }
+
+  /** A plan with no collation in its trait set is returned unchanged, with no spurious Sort. */
+  @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,30 @@ public void testComplexQuery() throws Exception {
     }
   }
 
+  @Test
+  public void testCompositeCollationQueryExecutes() throws Exception {
+    // `eval x = age` after `sort age` yields a top projection with a composite collation trait
+    // (x mirrors age). UnifiedQueryPlanner#preserveCollation used to crash building this plan;
+    // verify it now compiles and executes end to end.
+    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,33 @@ public void testMultiplePPLQueryExecutionWithSameContext() throws Exception {
     }
   }
 
+  @Test
+  public void testSortThenEvalCopyExecutesEndToEnd() throws Exception {
+    // `eval age_copy = age` after `sort age` makes the top projection carry more than one derived
+    // collation (age_copy mirrors age). UnifiedQueryPlanner#preserveCollation used to throw reading
+    // that composite trait; the query must now plan and execute against a real OpenSearch index.
+    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.