Support ORDER BY clauses in pagination and queries without FROM.

Signed-off-by: Yury-Fridlyand <[email protected]>

core/src/main/java/org/opensearch/sql/executor/pagination/CanPaginateVisitor.java

@@ -8,8 +8,12 @@
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.expression.AllFields;
+import org.opensearch.sql.ast.expression.QualifiedName;
+import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.Relation;
+import org.opensearch.sql.ast.tree.Sort;
+import org.opensearch.sql.ast.tree.Values;
 
 /**
  * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache.
@@ -19,9 +23,13 @@
  * Currently, the conditions are:
  * - only projection of a relation is supported.
  * - projection only has * (a.k.a. allFields).
- * - Relation only scans one table
+ * - Relation only scans one table.
  * - The table is an open search index.
- * So it accepts only queries like `select * from $index`
+ * - Query has <pre>ORDER BY</pre> clause, which refers columns only by name or by ordinal
+ *   (without any expression there).
+ * - Query has no FROM clause.
+ * So it accepts only queries like
+ * <pre>SELECT * FROM $index [ORDER BY $col1, [$col2]] [LIMIT n]</pre>
  * See PaginatedPlanCache.canConvertToCursor for usage.
  */
 public class CanPaginateVisitor extends AbstractNodeVisitor<Boolean, Object> {
@@ -36,16 +44,43 @@ public Boolean visitRelation(Relation node, Object context) {
     return Boolean.TRUE;
   }
 
+  protected Boolean canPaginate(Node node, Object context) {
+    var childList = node.getChild();
+    if (childList != null) {
+      return childList.stream().allMatch(n -> n.accept(this, context));
+    }
+    return Boolean.TRUE;
+  }
+
+  // Only column references in ORDER BY clause are supported in pagination,
+  // because expressions can't be pushed down due to #1471.
+  // https://github.com/opensearch-project/sql/issues/1471
+  @Override
+  public Boolean visitSort(Sort node, Object context) {
+    return node.getSortList().stream().allMatch(f -> f.getField() instanceof QualifiedName)
+        // TODO visit fields after #1500 merge https://github.com/opensearch-project/sql/pull/1500
+        && canPaginate(node, context);
+  }
+
+  // Queries without FROM clause are also supported
+  @Override
+  public Boolean visitValues(Values node, Object context) {
+    return Boolean.TRUE;
+  }
+
+  // Queries with LIMIT/OFFSET clauses are unsupported
+  @Override
+  public Boolean visitLimit(Limit node, Object context) {
+    return Boolean.FALSE;
+  }
+
   @Override
   public Boolean visitChildren(Node node, Object context) {
     return Boolean.FALSE;
   }
 
   @Override
   public Boolean visitProject(Project node, Object context) {
-    // Allow queries with 'SELECT *' only. Those restriction could be removed, but consider
-    // in-memory aggregation performed by window function (see WindowOperator).
-    // SELECT max(age) OVER (PARTITION BY city) ...
     var projections = node.getProjectList();
     if (projections.size() != 1) {
       return Boolean.FALSE;

core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java

@@ -14,6 +14,7 @@
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalNested;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
 import org.opensearch.sql.planner.logical.LogicalProject;
@@ -161,6 +162,12 @@ public PhysicalPlan visitCloseCursor(LogicalCloseCursor node, C context) {
     return new CursorCloseOperator(visitChild(node, context));
   }
 
+  // Called when paging query requested without `FROM` clause only
+  @Override
+  public PhysicalPlan visitPaginate(LogicalPaginate plan, C context) {
+    return visitChild(plan, context);
+  }
+
   protected PhysicalPlan visitChild(LogicalPlan node, C context) {
     // Logical operators visited here must have a single child
     return node.getChild().get(0).accept(this, context);

core/src/test/java/org/opensearch/sql/executor/pagination/CanPaginateVisitorTest.java

@@ -5,20 +5,60 @@
 
 package org.opensearch.sql.executor.pagination;
 
+import static org.junit.jupiter.api.Assertions.assertAll;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.withSettings;
+import static org.opensearch.sql.ast.dsl.AstDSL.agg;
+import static org.opensearch.sql.ast.dsl.AstDSL.aggregate;
+import static org.opensearch.sql.ast.dsl.AstDSL.alias;
+import static org.opensearch.sql.ast.dsl.AstDSL.allFields;
+import static org.opensearch.sql.ast.dsl.AstDSL.and;
+import static org.opensearch.sql.ast.dsl.AstDSL.argument;
+import static org.opensearch.sql.ast.dsl.AstDSL.between;
+import static org.opensearch.sql.ast.dsl.AstDSL.booleanLiteral;
+import static org.opensearch.sql.ast.dsl.AstDSL.caseWhen;
+import static org.opensearch.sql.ast.dsl.AstDSL.cast;
+import static org.opensearch.sql.ast.dsl.AstDSL.compare;
+import static org.opensearch.sql.ast.dsl.AstDSL.equalTo;
+import static org.opensearch.sql.ast.dsl.AstDSL.eval;
+import static org.opensearch.sql.ast.dsl.AstDSL.field;
+import static org.opensearch.sql.ast.dsl.AstDSL.filter;
+import static org.opensearch.sql.ast.dsl.AstDSL.function;
+import static org.opensearch.sql.ast.dsl.AstDSL.highlight;
+import static org.opensearch.sql.ast.dsl.AstDSL.in;
+import static org.opensearch.sql.ast.dsl.AstDSL.intLiteral;
+import static org.opensearch.sql.ast.dsl.AstDSL.intervalLiteral;
+import static org.opensearch.sql.ast.dsl.AstDSL.limit;
+import static org.opensearch.sql.ast.dsl.AstDSL.map;
+import static org.opensearch.sql.ast.dsl.AstDSL.not;
+import static org.opensearch.sql.ast.dsl.AstDSL.or;
+import static org.opensearch.sql.ast.dsl.AstDSL.project;
+import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
+import static org.opensearch.sql.ast.dsl.AstDSL.relation;
+import static org.opensearch.sql.ast.dsl.AstDSL.sort;
+import static org.opensearch.sql.ast.dsl.AstDSL.stringLiteral;
+import static org.opensearch.sql.ast.dsl.AstDSL.tableFunction;
+import static org.opensearch.sql.ast.dsl.AstDSL.unresolvedArg;
+import static org.opensearch.sql.ast.dsl.AstDSL.unresolvedAttr;
+import static org.opensearch.sql.ast.dsl.AstDSL.values;
+import static org.opensearch.sql.ast.dsl.AstDSL.when;
+import static org.opensearch.sql.ast.dsl.AstDSL.window;
+import static org.opensearch.sql.ast.dsl.AstDSL.xor;
 
 import java.util.List;
 import org.junit.jupiter.api.DisplayNameGeneration;
 import org.junit.jupiter.api.DisplayNameGenerator;
 import org.junit.jupiter.api.Test;
+import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.dsl.AstDSL;
+import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.Relation;
-import org.opensearch.sql.executor.pagination.CanPaginateVisitor;
 
 @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
 public class CanPaginateVisitorTest {
@@ -28,105 +68,164 @@ public class CanPaginateVisitorTest {
   @Test
   // select * from y
   public void accept_query_with_select_star_and_from() {
-    var plan = AstDSL.project(AstDSL.relation("dummy"), AstDSL.allFields());
+    var plan = project(relation("dummy"), allFields());
     assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select x from y
   public void reject_query_with_select_field_and_from() {
-    var plan = AstDSL.project(AstDSL.relation("dummy"), AstDSL.field("pewpew"));
+    var plan = project(relation("dummy"), field("pewpew"));
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select x,z from y
   public void reject_query_with_select_fields_and_from() {
-    var plan = AstDSL.project(AstDSL.relation("dummy"),
-        AstDSL.field("pewpew"), AstDSL.field("pewpew"));
+    var plan = project(relation("dummy"),
+        field("pewpew"), field("pewpew"));
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select x
-  public void reject_query_without_from() {
-    var plan = AstDSL.project(AstDSL.values(List.of(AstDSL.intLiteral(1))),
-        AstDSL.alias("1",AstDSL.intLiteral(1)));
-    assertFalse(plan.accept(visitor, null));
+  public void allow_query_without_from() {
+    // bypass check for `AllFields` in `visitProject`
+    // TODO remove after #1500 merge https://github.com/opensearch-project/sql/pull/1500
+    var visitor = new CanPaginateVisitor() {
+      @Override
+      public Boolean visitProject(Project node, Object context) {
+        return node.getChild().get(0).accept(this, context);
+      }
+    };
+    var plan = project(values(List.of(intLiteral(1))),
+        alias("1", intLiteral(1)));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select * from y limit z
-  public void reject_query_with_limit() {
-    var plan = AstDSL.project(AstDSL.limit(AstDSL.relation("dummy"), 1, 2), AstDSL.allFields());
+  public void reject_query_with_limit_but_no_offset() {
+    var plan = project(limit(relation("dummy"), 1, 0), allFields());
+    assertFalse(plan.accept(visitor, null));
+  }
+
+  @Test
+  // select * from y limit z, n
+  public void reject_query_with_offset() {
+    var plan = project(limit(relation("dummy"), 0, 1), allFields());
+    assertFalse(plan.accept(visitor, null));
+  }
+
+  // test added for coverage only
+  @Test
+  public void visitLimit() {
+    var visitor = new CanPaginateVisitor() {
+      @Override
+      public Boolean visitRelation(Relation node, Object context) {
+        return Boolean.FALSE;
+      }
+    };
+    var plan = project(limit(relation("dummy"), 0, 0), allFields());
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select * from y where z
   public void reject_query_with_where() {
-    var plan = AstDSL.project(AstDSL.filter(AstDSL.relation("dummy"),
-        AstDSL.booleanLiteral(true)), AstDSL.allFields());
+    var plan = project(filter(relation("dummy"),
+        booleanLiteral(true)), allFields());
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select * from y order by z
-  public void reject_query_with_order_by() {
-    var plan = AstDSL.project(AstDSL.sort(AstDSL.relation("dummy"), AstDSL.field("1")),
-        AstDSL.allFields());
+  public void allow_query_with_order_by_with_column_references_only() {
+    var plan = project(sort(relation("dummy"), field("1")), allFields());
+    assertTrue(plan.accept(visitor, null));
+  }
+
+  @Test
+  // select * from y order by func(z)
+  public void reject_query_with_order_by_with_an_expression() {
+    var plan = project(sort(relation("dummy"), field(function("func"))),
+        allFields());
+    assertFalse(plan.accept(visitor, null));
+  }
+
+  // test added for coverage only
+  @Test
+  public void visitSort() {
+    var visitor = new CanPaginateVisitor() {
+      @Override
+      public Boolean visitRelation(Relation node, Object context) {
+        return Boolean.FALSE;
+      }
+    };
+    var plan = project(sort(relation("dummy"), field("1")), allFields());
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select * from y group by z
   public void reject_query_with_group_by() {
-    var plan = AstDSL.project(AstDSL.agg(
-        AstDSL.relation("dummy"), List.of(), List.of(), List.of(AstDSL.field("1")), List.of()),
-        AstDSL.allFields());
+    var plan = project(agg(
+        relation("dummy"), List.of(), List.of(), List.of(field("1")), List.of()),
+        allFields());
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select agg(x) from y
   public void reject_query_with_aggregation_function() {
-    var plan = AstDSL.project(AstDSL.agg(
-        AstDSL.relation("dummy"),
-        List.of(AstDSL.alias("agg", AstDSL.aggregate("func", AstDSL.field("pewpew")))),
+    var plan = project(agg(
+        relation("dummy"),
+        List.of(alias("agg", aggregate("func", field("pewpew")))),
         List.of(), List.of(), List.of()),
-        AstDSL.allFields());
+        allFields());
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select window(x) from y
   public void reject_query_with_window_function() {
-    var plan = AstDSL.project(AstDSL.relation("dummy"),
-        AstDSL.alias("pewpew",
-            AstDSL.window(
-                AstDSL.aggregate("func", AstDSL.field("pewpew")),
-                    List.of(AstDSL.qualifiedName("1")), List.of())));
+    var plan = project(relation("dummy"),
+        alias("pewpew",
+            window(
+                aggregate("func", field("pewpew")),
+                    List.of(qualifiedName("1")), List.of())));
     assertFalse(plan.accept(visitor, null));
   }
 
   @Test
   // select * from y, z
   public void reject_query_with_select_from_multiple_indices() {
     var plan = mock(Project.class);
-    when(plan.getChild()).thenReturn(List.of(AstDSL.relation("dummy"), AstDSL.relation("pummy")));
-    when(plan.getProjectList()).thenReturn(List.of(AstDSL.allFields()));
+    when(plan.getChild()).thenReturn(List.of(relation("dummy"), relation("pummy")));
+    when(plan.getProjectList()).thenReturn(List.of(allFields()));
     assertFalse(visitor.visitProject(plan, null));
   }
 
   @Test
   // unreal case, added for coverage only
   public void reject_project_when_relation_has_child() {
-    var relation = mock(Relation.class, withSettings().useConstructor(AstDSL.qualifiedName("42")));
-    when(relation.getChild()).thenReturn(List.of(AstDSL.relation("pewpew")));
+    var relation = mock(Relation.class, withSettings().useConstructor(qualifiedName("42")));
+    when(relation.getChild()).thenReturn(List.of(relation("pewpew")));
     when(relation.accept(visitor, null)).thenCallRealMethod();
     var plan = mock(Project.class);
     when(plan.getChild()).thenReturn(List.of(relation));
-    when(plan.getProjectList()).thenReturn(List.of(AstDSL.allFields()));
+    when(plan.getProjectList()).thenReturn(List.of(allFields()));
     assertFalse(visitor.visitProject((Project) plan, null));
   }
+
+  @Test
+  // test added just for coverage
+  public void canPaginate() {
+    assertTrue(visitor.canPaginate(new Node() {
+      @Override
+      public List<? extends Node> getChild() {
+        return null;
+      }
+    }, null));
+  }
 }

core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java

@@ -58,12 +58,17 @@
 import org.opensearch.sql.expression.window.WindowDefinition;
 import org.opensearch.sql.expression.window.ranking.RowNumberFunction;
 import org.opensearch.sql.planner.logical.LogicalCloseCursor;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanDSL;
+import org.opensearch.sql.planner.logical.LogicalProject;
 import org.opensearch.sql.planner.logical.LogicalRelation;
+import org.opensearch.sql.planner.logical.LogicalValues;
 import org.opensearch.sql.planner.physical.CursorCloseOperator;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlanDSL;
+import org.opensearch.sql.planner.physical.ProjectOperator;
+import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.storage.StorageEngine;
 import org.opensearch.sql.storage.Table;
 import org.opensearch.sql.storage.TableScanOperator;
@@ -273,4 +278,14 @@ public void visitCloseCursor_should_build_CursorCloseOperator() {
     assertTrue(implemented instanceof CursorCloseOperator);
     assertSame(physicalChild, implemented.getChild().get(0));
   }
+
+  @Test
+  public void visitPaginate_should_remove_it_from_tree() {
+    var logicalPlanTree = new LogicalPaginate(42, List.of(
+        new LogicalProject(
+            new LogicalValues(List.of(List.of())), List.of(), List.of())));
+    var physicalPlanTree = new ProjectOperator(
+        new ValuesOperator(List.of(List.of())), List.of(), List.of());
+    assertEquals(physicalPlanTree, logicalPlanTree.accept(implementor, null));
+  }
 }