Fix #5114: preserve head/TopK semantics for sort-expression pushdown

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -2033,6 +2033,17 @@ public void testTopKThenSortExplain() throws IOException {
                 + "| fields age"));
   }
 
+  @Test
+  public void testIssue5114SortExprHeadExplain() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    String query =
+        "source=opensearch-sql_test_index_account | eval a = rand() | sort a | fields"
+            + " account_number | head 5";
+    var result = explainQueryYaml(query);
+    String expected = loadExpectedPlan("explain_issue_5114_sort_expr_head_push.yaml");
+    assertYamlEqualsIgnoreId(expected, result);
+  }
+
   @Test
   public void testGeoIpPushedInAgg() throws IOException {
     // This explain IT verifies that externally registered UDF can be properly pushed down

integ-test/src/test/resources/expectedOutput/calcite/explain_issue_5114_sort_expr_head_push.yaml

@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0])
+        LogicalSort(sort0=[$17], dir0=[ASC-nulls-first], fetch=[5])
+          LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], a=[RAND()])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number], SORT_EXPR->[RAND() ASCENDING NULLS_FIRST], LIMIT->5, LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["account_number"],"excludes":[]},"sort":[{"_script":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQAbnsKICAib3AiOiB7CiAgICAibmFtZSI6ICJSQU5EIiwKICAgICJraW5kIjogIk9USEVSX0ZVTkNUSU9OIiwKICAgICJzeW50YXgiOiAiRlVOQ1RJT04iCiAgfSwKICAib3BlcmFuZHMiOiBbXQp9\"}","lang":"opensearch_compounded_script","params":{"MISSING_MAX":false,"utcTimestamp": 0,"SOURCES":[],"DIGESTS":[]}},"type":"number","order":"asc"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])

integ-test/src/yamlRestTest/resources/rest-api-spec/test/issues/5114.yml

@@ -0,0 +1,97 @@
+setup:
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled: true
+
+  - do:
+      indices.create:
+        index: issue5114
+        body:
+          settings:
+            number_of_shards: 1
+            number_of_replicas: 0
+          mappings:
+            properties:
+              id:
+                type: integer
+              name:
+                type: keyword
+              reportsTo:
+                type: keyword
+
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{"index": {"_index": "issue5114", "_id": "1"}}'
+          - '{"id": 1, "name": "Dev", "reportsTo": "Eliot"}'
+          - '{"index": {"_index": "issue5114", "_id": "2"}}'
+          - '{"id": 2, "name": "Eliot", "reportsTo": "Ron"}'
+          - '{"index": {"_index": "issue5114", "_id": "3"}}'
+          - '{"id": 3, "name": "Ron", "reportsTo": "Andrew"}'
+          - '{"index": {"_index": "issue5114", "_id": "4"}}'
+          - '{"id": 4, "name": "Andrew", "reportsTo": null}'
+          - '{"index": {"_index": "issue5114", "_id": "5"}}'
+          - '{"id": 5, "name": "Asya", "reportsTo": "Ron"}'
+          - '{"index": {"_index": "issue5114", "_id": "6"}}'
+          - '{"id": 6, "name": "Dan", "reportsTo": "Andrew"}'
+
+---
+teardown:
+  - do:
+      indices.delete:
+        index: issue5114
+        ignore_unavailable: true
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled: false
+
+---
+"Issue 5114: head should be preserved for non-order-equivalent deterministic sort expression":
+  - skip:
+      features:
+        - headers
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=issue5114 | eval a = abs(id) + 1 | sort a | fields id | head 5
+
+  - match: { total: 5 }
+  - match: { schema: [ { name: id, type: int } ] }
+
+---
+"Issue 5114: head should be preserved for non-deterministic sort expression":
+  - skip:
+      features:
+        - headers
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=issue5114 | eval a = rand() | sort a | fields id | head 5
+
+  - match: { total: 5 }
+  - match: { schema: [ { name: id, type: int } ] }
+
+---
+"Issue 5114 control: order-equivalent expression remains correct":
+  - skip:
+      features:
+        - headers
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=issue5114 | eval a = id + 1 | sort a | fields id | head 5
+
+  - match: { total: 5 }
+  - match: { schema: [ { name: id, type: int } ] }
+  - match: { datarows: [ [ 1 ], [ 2 ], [ 3 ], [ 4 ], [ 5 ] ] }

opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/SortExprIndexScanRule.java

@@ -11,6 +11,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Predicate;
+import org.apache.calcite.adapter.enumerable.EnumerableLimitSort;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelFieldCollation.Direction;
@@ -47,6 +48,11 @@ protected SortExprIndexScanRule(SortExprIndexScanRule.Config config) {
   @Override
   protected void onMatchImpl(RelOptRuleCall call) {
     final Sort sort = call.rel(0);
+    // EnumerableLimitSort carries fetch semantics; this rule doesn't preserve it on physical
+    // scans because limit pushdown path is logical-only.
+    if (sort instanceof EnumerableLimitSort) {
+      return;
+    }
     final Project project = call.rel(1);
     final AbstractCalciteIndexScan scan = call.rel(2);