Push down sort through eval

core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java

@@ -48,6 +48,7 @@ public static LogicalPlanOptimizer create() {
             new MergeFilterAndFilter(),
             new PushFilterUnderSort(),
             EvalPushDown.PUSH_DOWN_LIMIT,
+            EvalPushDown.PUSH_DOWN_SORT,
             /*
              * Phase 2: Transformations that rely on data source push down capability
              */

core/src/main/java/org/opensearch/sql/planner/optimizer/rule/EvalPushDown.java

@@ -7,20 +7,29 @@
 
 import static org.opensearch.sql.planner.optimizer.pattern.Patterns.evalCapture;
 import static org.opensearch.sql.planner.optimizer.pattern.Patterns.limit;
+import static org.opensearch.sql.planner.optimizer.pattern.Patterns.sort;
 import static org.opensearch.sql.planner.optimizer.rule.EvalPushDown.EvalPushDownBuilder.match;
 
 import com.facebook.presto.matching.Capture;
 import com.facebook.presto.matching.Captures;
 import com.facebook.presto.matching.Pattern;
 import com.facebook.presto.matching.pattern.CapturePattern;
 import com.facebook.presto.matching.pattern.WithPattern;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
+import java.util.stream.Collectors;
 import lombok.Getter;
 import lombok.experimental.Accessors;
+import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.sql.ast.tree.Sort.SortOption;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.planner.logical.LogicalEval;
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.planner.logical.LogicalSort;
 import org.opensearch.sql.planner.optimizer.Rule;
 
 /**
@@ -42,6 +51,38 @@ public class EvalPushDown<T extends LogicalPlan> implements Rule<T> {
                 return logicalEval;
               });
 
+  public static final Rule<LogicalSort> PUSH_DOWN_SORT =
+      match(sort(evalCapture()))
+          .apply(
+              (sort, logicalEval) -> {
+                List<LogicalPlan> child = logicalEval.getChild();
+                Map<ReferenceExpression, Expression> evalExpressionMap =
+                    logicalEval.getExpressions().stream()
+                        .collect(Collectors.toMap(Pair::getKey, Pair::getValue));
+                List<Pair<SortOption, Expression>> sortList = sort.getSortList();
+                List<Pair<SortOption, Expression>> newSortList = new ArrayList<>();
+                for (Pair<SortOption, Expression> pair : sortList) {
+                  /*
+                   Narrow down the optimization to only support:
+                   1. The expression in sort and replaced expression are both ReferenceExpression.
+                   2. No internal reference in eval.
+                  */
+                  if (pair.getRight() instanceof ReferenceExpression) {
+                    ReferenceExpression ref = (ReferenceExpression) pair.getRight();
+                    Expression replacedExpr = evalExpressionMap.getOrDefault(ref, ref);
+                    if (replacedExpr instanceof ReferenceExpression) {
+                      ReferenceExpression newRef = (ReferenceExpression) replacedExpr;
+                      if (!evalExpressionMap.containsKey(newRef)) {
+                        newSortList.add(Pair.of(pair.getLeft(), newRef));
+                      } else return sort;
+                    } else return sort;
+                  } else return sort;
+                }
+                sort = new LogicalSort(child.getFirst(), newSortList);
+                logicalEval.replaceChildPlans(List.of(sort));
+                return logicalEval;
+              });
+
   private final Capture<LogicalEval> capture;
 
   @Accessors(fluent = true)

core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java

@@ -42,6 +42,7 @@
 import org.mockito.Spy;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.ast.tree.Sort;
+import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
@@ -368,6 +369,48 @@ void push_limit_through_eval_into_scan() {
         optimize(limit(eval(relation("schema", table), evalExpr), 10, 5)));
   }
 
+  /** Sort - Eval --> Eval - Sort. */
+  @Test
+  void push_sort_under_eval() {
+    ReferenceExpression sortRef = DSL.ref("intV", INTEGER);
+    ReferenceExpression evalRef = DSL.ref("name1", INTEGER);
+    Pair<ReferenceExpression, Expression> evalExpr = Pair.of(evalRef, DSL.ref("name", STRING));
+    Pair<SortOption, Expression> sortExpr = Pair.of(Sort.SortOption.DEFAULT_ASC, sortRef);
+    assertEquals(
+        eval(sort(tableScanBuilder, sortExpr), evalExpr),
+        optimize(sort(eval(relation("schema", table), evalExpr), sortExpr)));
+
+    // don't push sort if sort field is not ReferenceExpression
+    Expression nonRefExpr = DSL.add(DSL.ref("intV", INTEGER), DSL.literal(1));
+    Pair<SortOption, Expression> sortExprWithNonRef =
+        Pair.of(Sort.SortOption.DEFAULT_ASC, nonRefExpr);
+    LogicalPlan originPlan = sort(eval(relation("schema", table), evalExpr), sortExprWithNonRef);
+    assertEquals(originPlan, optimize(originPlan));
+
+    // don't push sort if replaced expr in eval is not ReferenceExpression
+    Pair<ReferenceExpression, Expression> evalExprWithNonRef = Pair.of(sortRef, nonRefExpr);
+    originPlan = sort(eval(relation("schema", table), evalExprWithNonRef), sortExpr);
+    assertEquals(originPlan, optimize(originPlan));
+
+    // don't push sort if there are internal reference in eval
+    Pair<ReferenceExpression, Expression> evalExpr2 = Pair.of(sortRef, evalRef);
+    originPlan = sort(eval(relation("schema", table), evalExpr, evalExpr2), sortExpr);
+    assertEquals(originPlan, optimize(originPlan));
+  }
+
+  /** Sort - Eval - Scan --> Eval - Scan. */
+  @Test
+  void push_sort_through_eval_into_scan() {
+    when(tableScanBuilder.pushDownSort(any())).thenReturn(true);
+    Pair<ReferenceExpression, Expression> evalExpr =
+        Pair.of(DSL.ref("name1", STRING), DSL.ref("name", STRING));
+    Pair<SortOption, Expression> sortExpr =
+        Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("intV", INTEGER));
+    assertEquals(
+        eval(tableScanBuilder, evalExpr),
+        optimize(sort(eval(relation("schema", table), evalExpr), sortExpr)));
+  }
+
   private LogicalPlan optimize(LogicalPlan plan) {
     final LogicalPlanOptimizer optimizer = LogicalPlanOptimizer.create();
     return optimizer.optimize(plan);

integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java

@@ -89,6 +89,19 @@ public void testLimitPushDownExplain() throws Exception {
                 + "| fields ageMinus"));
   }
 
+  @Test
+  public void testSortPushDownThroughEvalExplain() throws Exception {
+    String expected = loadFromFile("expectedOutput/ppl/explain_sort_push_through_eval.json");
+
+    assertJsonEquals(
+        expected,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| eval newAge = age"
+                + "| sort newAge"
+                + "| fields newAge"));
+  }
+
   String loadFromFile(String filename) throws Exception {
     URI uri = Resources.getResource(filename).toURI();
     return new String(Files.readAllBytes(Paths.get(uri)));

integ-test/src/test/resources/expectedOutput/ppl/explain_sort_push_through_eval.json

@@ -0,0 +1,27 @@
+{
+  "root": {
+    "name": "ProjectOperator",
+    "description": {
+      "fields": "[newAge]"
+    },
+    "children": [
+      {
+        "name": "EvalOperator",
+        "description": {
+          "expressions": {
+            "newAge": "age"
+          }
+        },
+        "children": [
+          {
+            "name": "OpenSearchIndexScan",
+            "description": {
+              "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, searchDone=false)"
+            },
+            "children": []
+          }
+        ]
+      }
+    ]
+  }
+}