Pagination Phase 2: Support ORDER BY clauses and queries without `F…

…ROM`. (#1599) * Support `ORDER BY` clauses in pagination and queries without `FROM`. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix IT. Signed-off-by: Yury-Fridlyand <[email protected]> --------- Signed-off-by: Yury-Fridlyand <[email protected]> (cherry picked from commit 94d5479)
opensearch-project · Jun 15, 2023 · 29344e5 · 29344e5
1 parent d2a82ae
commit 29344e5
Show file tree

Hide file tree

Showing 8 changed files with 217 additions and 52 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/executor/pagination/CanPaginateVisitor.java b/core/src/main/java/org/opensearch/sql/executor/pagination/CanPaginateVisitor.java
@@ -48,7 +48,6 @@
  * Currently, V2 engine does not support queries with:
  * - aggregation (GROUP BY clause or aggregation functions like min/max)
  * - in memory aggregation (window function)
- * - ORDER BY clause
  * - LIMIT/OFFSET clause(s)
  * - without FROM clause
  * - JOIN
@@ -68,14 +67,24 @@ public Boolean visitRelation(Relation node, Object context) {
     return Boolean.TRUE;
   }
 
-  private Boolean canPaginate(Node node, Object context) {
+  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
+            && visitField(f, context))
+        && canPaginate(node, context);
+  }
+
   // For queries with WHERE clause:
   @Override
   public Boolean visitFilter(Filter node, Object context) {
@@ -88,19 +97,13 @@ public Boolean visitAggregation(Aggregation node, Object context) {
     return Boolean.FALSE;
   }
 
-  // Queries with ORDER BY clause are not supported
-  @Override
-  public Boolean visitSort(Sort node, Object context) {
-    return Boolean.FALSE;
-  }
-
-  // Queries without FROM clause are not supported
+  // For queries without FROM clause:
   @Override
   public Boolean visitValues(Values node, Object context) {
-    return Boolean.FALSE;
+    return Boolean.TRUE;
   }
 
-  // Queries with LIMIT clause are not supported
+  // Queries with LIMIT/OFFSET clauses are unsupported
   @Override
   public Boolean visitLimit(Limit node, Object context) {
     return Boolean.FALSE;

diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/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);

diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/CanPaginateVisitorTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/CanPaginateVisitorTest.java
@@ -57,6 +57,7 @@
 import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.Argument;
 import org.opensearch.sql.ast.expression.DataType;
+import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
@@ -119,10 +120,10 @@ public void allow_query_with_select_fields_and_from() {
 
   @Test
   // select x
-  public void reject_query_without_from() {
+  public void allow_query_without_from() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", intLiteral(1)));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
@@ -165,63 +166,63 @@ public List<? extends Node> getChild() {
   public void visitEqualTo() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", equalTo(intLiteral(1), intLiteral(1))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select interval
   public void visitInterval() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", intervalLiteral(intLiteral(1), DataType.INTEGER, "days")));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select a != b
   public void visitCompare() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", compare("!=", intLiteral(1), intLiteral(1))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select NOT a
   public void visitNot() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", not(booleanLiteral(true))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select a OR b
   public void visitOr() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", or(booleanLiteral(true), booleanLiteral(false))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select a AND b
   public void visitAnd() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", and(booleanLiteral(true), booleanLiteral(false))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select a XOR b
   public void visitXor() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", xor(booleanLiteral(true), booleanLiteral(false))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select f()
   public void visitFunction() {
     var plan = project(values(List.of(intLiteral(1))),
         function("func"));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
@@ -237,7 +238,7 @@ public void visitNested() {
   public void visitIn() {
     // test combinations of acceptable and not acceptable args for coverage
     assertAll(
-        () -> assertFalse(project(values(List.of(intLiteral(1))), alias("1", in(field("a"))))
+        () -> assertTrue(project(values(List.of(intLiteral(1))), alias("1", in(field("a"))))
             .accept(visitor, null)),
         () -> assertFalse(project(values(List.of(intLiteral(1))),
                 alias("1", in(field("a"), map("1", "2"))))
@@ -253,23 +254,23 @@ public void visitIn() {
   public void visitBetween() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", between(field("a"), intLiteral(1), intLiteral(2))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select a CASE 1 WHEN 2
   public void visitCase() {
     var plan = project(values(List.of(intLiteral(1))),
         alias("1", caseWhen(intLiteral(1), when(intLiteral(3), intLiteral(4)))));
-    assertFalse(plan.accept(visitor, null));
+    assertTrue(plan.accept(visitor, null));
   }
 
   @Test
   // select CAST(a as TYPE)
   public void visitCast() {
     // test combinations of acceptable and not acceptable args for coverage
     assertAll(
-        () -> assertFalse(project(values(List.of(intLiteral(1))),
+        () -> assertTrue(project(values(List.of(intLiteral(1))),
                 alias("1", cast(intLiteral(2), stringLiteral("int"))))
             .accept(visitor, null)),
         () -> assertFalse(project(values(List.of(intLiteral(1))),
@@ -323,8 +324,28 @@ public void accept_query_with_highlight_and_relevance_func() {
 
   @Test
   // select * from y limit z
-  public void reject_query_with_limit() {
-    var plan = project(limit(relation("dummy"), 1, 2), 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));
   }
 
@@ -338,12 +359,40 @@ public void allow_query_with_where() {
 
   @Test
   // select * from y order by z
-  public void reject_query_with_order_by() {
-    var plan = project(sort(relation("dummy"), field("1")),
+  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() {
+    CanPaginateVisitor 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));
+    visitor = new CanPaginateVisitor() {
+      @Override
+      public Boolean visitField(Field node, Object context) {
+        return Boolean.FALSE;
+      }
+    };
+    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() {
@@ -396,22 +445,6 @@ public void reject_project_when_relation_has_child() {
     assertFalse(visitor.visitProject((Project) plan, null));
   }
 
-  @Test
-  // test combinations of acceptable and not acceptable args for coverage
-  public void canPaginate() {
-    assertAll(
-        () -> assertFalse(project(values(List.of(intLiteral(1))),
-                function("func", intLiteral(1), intLiteral(1)))
-            .accept(visitor, null)),
-        () -> assertFalse(project(values(List.of(intLiteral(1))),
-                function("func", intLiteral(1), map("1", "2")))
-            .accept(visitor, null)),
-        () -> assertFalse(project(values(List.of(intLiteral(1))),
-                function("func", map("1", "2"), intLiteral(1)))
-            .accept(visitor, null))
-    );
-  }
-
   @Test
   // test combinations of acceptable and not acceptable args for coverage
   public void visitFilter() {

diff --git a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java b/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));
+  }
 }
diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTestBase.java b/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTestBase.java
@@ -8,6 +8,9 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.io.IOException;
+import java.io.ObjectInput;
+import java.io.ObjectOutput;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -22,6 +25,7 @@
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.env.Environment;
+import org.opensearch.sql.planner.SerializablePlan;
 
 public class PhysicalPlanTestBase {
 
@@ -208,7 +212,7 @@ protected static PhysicalPlan testScan(List<ExprValue> inputs) {
     return new TestScan(inputs);
   }
 
-  protected static class TestScan extends PhysicalPlan {
+  protected static class TestScan extends PhysicalPlan implements SerializablePlan {
     private final Iterator<ExprValue> iterator;
 
     public TestScan() {
@@ -238,5 +242,21 @@ public boolean hasNext() {
     public ExprValue next() {
       return iterator.next();
     }
+
+    @Override
+    public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
+    }
+
+    @Override
+    public void writeExternal(ObjectOutput out) throws IOException {
+    }
+
+    public boolean equals(final Object o) {
+      return o == this || o.hashCode() == hashCode();
+    }
+
+    public int hashCode() {
+      return 42;
+    }
   }
 }