add visitTakeOrdered

Signed-off-by: Heng Qian <[email protected]>
opensearch-project · Jul 26, 2024 · b03cfc0 · b03cfc0
1 parent c940a0d
commit b03cfc0
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 1 deletion.
diff --git a/core/src/main/java/org/opensearch/sql/executor/Explain.java b/core/src/main/java/org/opensearch/sql/executor/Explain.java
@@ -30,6 +30,7 @@
 import org.opensearch.sql.planner.physical.RemoveOperator;
 import org.opensearch.sql.planner.physical.RenameOperator;
 import org.opensearch.sql.planner.physical.SortOperator;
+import org.opensearch.sql.planner.physical.TakeOrderedOperator;
 import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.planner.physical.WindowOperator;
 import org.opensearch.sql.storage.TableScanOperator;
@@ -73,6 +74,19 @@ public ExplainResponseNode visitSort(SortOperator node, Object context) {
                 ImmutableMap.of("sortList", describeSortList(node.getSortList()))));
   }
 
+  @Override
+  public ExplainResponseNode visitTakeOrdered(TakeOrderedOperator node, Object context) {
+    return explain(
+        node,
+        context,
+        explainNode ->
+            explainNode.setDescription(
+                ImmutableMap.of(
+                    "limit", node.getLimit(),
+                    "offset", node.getOffset(),
+                    "sortList", describeSortList(node.getSortList()))));
+  }
+
   @Override
   public ExplainResponseNode visitTableScan(TableScanOperator node, Object context) {
     return explain(

diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java
@@ -72,6 +72,10 @@ public R visitSort(SortOperator node, C context) {
     return visitNode(node, context);
   }
 
+  public R visitTakeOrdered(TakeOrderedOperator node, C context) {
+    return visitNode(node, context);
+  }
+
   public R visitRareTopN(RareTopNOperator node, C context) {
     return visitNode(node, context);
   }

diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/TakeOrderedOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/TakeOrderedOperator.java
@@ -81,7 +81,7 @@ public TakeOrderedOperator(
 
   @Override
   public <R, C> R accept(PhysicalPlanNodeVisitor<R, C> visitor, C context) {
-    return visitor.visitNode(this, context);
+    return visitor.visitTakeOrdered(this, context);
   }
 
   @Override

diff --git a/core/src/test/java/org/opensearch/sql/executor/ExplainTest.java b/core/src/test/java/org/opensearch/sql/executor/ExplainTest.java
@@ -27,6 +27,7 @@
 import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.remove;
 import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.rename;
 import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.sort;
+import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.takeOrdered;
 import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.values;
 import static org.opensearch.sql.planner.physical.PhysicalPlanDSL.window;
 
@@ -220,6 +221,26 @@ void can_explain_limit() {
         explain.apply(plan));
   }
 
+  @Test
+  void can_explain_takeOrdered() {
+    Pair<Sort.SortOption, Expression> sort =
+        ImmutablePair.of(Sort.SortOption.DEFAULT_ASC, ref("a", INTEGER));
+    PhysicalPlan plan = takeOrdered(tableScan, 10, 5, sort);
+    assertEquals(
+        new ExplainResponse(
+            new ExplainResponseNode(
+                "TakeOrderedOperator",
+                Map.of(
+                    "limit",
+                    10,
+                    "offset",
+                    5,
+                    "sortList",
+                    Map.of("a", Map.of("sortOrder", "ASC", "nullOrder", "NULL_FIRST"))),
+                singletonList(tableScan.explainNode()))),
+        explain.apply(plan));
+  }
+
   @Test
   void can_explain_nested() {
     Set<String> nestedOperatorArgs = Set.of("message.info", "message");

diff --git a/...n/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java b/...n/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java
@@ -23,6 +23,7 @@
 import org.opensearch.sql.planner.physical.RemoveOperator;
 import org.opensearch.sql.planner.physical.RenameOperator;
 import org.opensearch.sql.planner.physical.SortOperator;
+import org.opensearch.sql.planner.physical.TakeOrderedOperator;
 import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.planner.physical.WindowOperator;
 import org.opensearch.sql.storage.TableScanOperator;
@@ -130,6 +131,17 @@ public PhysicalPlan visitSort(SortOperator node, Object context) {
     return doProtect(new SortOperator(visitInput(node.getInput(), context), node.getSortList()));
   }
 
+  /** Decorate with {@link ResourceMonitorPlan}. */
+  @Override
+  public PhysicalPlan visitTakeOrdered(TakeOrderedOperator node, Object context) {
+    return doProtect(
+        new TakeOrderedOperator(
+            visitInput(node.getInput(), context),
+            node.getLimit(),
+            node.getOffset(),
+            node.getSortList()));
+  }
+
   /**
    * Values are a sequence of rows of literal value in memory which doesn't need memory protection.
    */

diff --git a/...va/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/...va/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java
@@ -306,6 +306,15 @@ void do_nothing_with_CursorCloseOperator_and_children() {
     verify(child, never()).accept(executionProtector, null);
   }
 
+  @Test
+  public void test_visitTakeOrdered() {
+    Pair<Sort.SortOption, Expression> sort =
+        ImmutablePair.of(Sort.SortOption.DEFAULT_ASC, ref("a", INTEGER));
+    PhysicalPlan physicalPlanTree =
+        PhysicalPlanDSL.takeOrdered(PhysicalPlanDSL.values(emptyList()), 10, 5, sort);
+    assertEquals(resourceMonitor(physicalPlanTree), executionProtector.doProtect(physicalPlanTree));
+  }
+
   PhysicalPlan resourceMonitor(PhysicalPlan input) {
     return new ResourceMonitorPlan(input, resourceMonitor);
   }