diff --git a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java b/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java index 9440833503f..bc86ac430de 100644 --- a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java +++ b/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 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; } diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java index 008121e8377..085171b7764 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java +++ b/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"); + } } diff --git a/api/src/test/java/org/opensearch/sql/api/compiler/UnifiedQueryCompilerTest.java b/api/src/test/java/org/opensearch/sql/api/compiler/UnifiedQueryCompilerTest.java index f344969cdc3..9ca25f1e523 100644 --- a/api/src/test/java/org/opensearch/sql/api/compiler/UnifiedQueryCompilerTest.java +++ b/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); diff --git a/integ-test/src/test/java/org/opensearch/sql/api/UnifiedQueryOpenSearchIT.java b/integ-test/src/test/java/org/opensearch/sql/api/UnifiedQueryOpenSearchIT.java index 7f7d31790f1..df7557976f0 100644 --- a/integ-test/src/test/java/org/opensearch/sql/api/UnifiedQueryOpenSearchIT.java +++ b/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.