Merge paginated plan optimizer into the regular optimizer.

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -15,9 +15,7 @@
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.planner.PlanContext;
 import org.opensearch.sql.planner.Planner;
-import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
-import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 /**
@@ -30,15 +28,7 @@ public class QueryService {
 
   private final ExecutionEngine executionEngine;
 
-  /**
-   * There are two planners, one - to handle pagination requests (cursor/scroll) only and
-   * another one for everything else.
-   * @see OpenSearchPluginModule#queryPlanFactory (:plugin module)
-   * @see LogicalPlanOptimizer#paginationCreate
-   * @see QueryService
-   */
   private final Planner planner;
-  private final Planner paginationPlanner;
 
   /**
    * Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.
@@ -115,6 +105,6 @@ public LogicalPlan analyze(UnresolvedPlan plan) {
    * Translate {@link LogicalPlan} to {@link PhysicalPlan}.
    */
   public PhysicalPlan plan(LogicalPlan plan) {
-    return plan instanceof LogicalPaginate ? paginationPlanner.plan(plan) : planner.plan(plan);
+    return planner.plan(plan);
   }
 }

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java

@@ -54,6 +54,10 @@ public static LogicalPlan rename(
     return new LogicalRename(input, renameMap);
   }
 
+  public  static LogicalPlan paginate(LogicalPlan input, int fetchSize) {
+    return new LogicalPaginate(fetchSize, List.of(input));
+  }
+
   public static LogicalPlan project(LogicalPlan input, NamedExpression... fields) {
     return new LogicalProject(input, Arrays.asList(fields), ImmutableList.of());
   }

core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java

@@ -9,7 +9,6 @@
 import com.google.common.collect.ImmutableList;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.Setter;
 import lombok.ToString;
 import org.opensearch.sql.storage.Table;
 
@@ -26,18 +25,13 @@ public class LogicalRelation extends LogicalPlan {
   @Getter
   private final Table table;
 
-  @Getter
-  @Setter
-  private Integer pageSize;
-
   /**
    * Constructor of LogicalRelation.
    */
   public LogicalRelation(String relationName, Table table) {
     super(ImmutableList.of());
     this.relationName = relationName;
     this.table = table;
-    this.pageSize = null;
   }
 
   @Override

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

@@ -16,7 +16,6 @@
 import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder;
 import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter;
 import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
-import org.opensearch.sql.planner.optimizer.rule.PushPageSize;
 import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
 import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown;
 import org.opensearch.sql.planner.optimizer.rule.write.CreateTableWriteBuilder;
@@ -53,29 +52,6 @@ public static LogicalPlanOptimizer create() {
          * Phase 2: Transformations that rely on data source push down capability
          */
         new CreateTableScanBuilder(),
-        TableScanPushDown.PUSH_DOWN_FILTER,
-        TableScanPushDown.PUSH_DOWN_AGGREGATION,
-        TableScanPushDown.PUSH_DOWN_SORT,
-        TableScanPushDown.PUSH_DOWN_LIMIT,
-        TableScanPushDown.PUSH_DOWN_HIGHLIGHT,
-        TableScanPushDown.PUSH_DOWN_PROJECT,
-        new CreateTableWriteBuilder()));
-  }
-
-  /**
-   * Create {@link LogicalPlanOptimizer} with pre-defined rules.
-   */
-  public static LogicalPlanOptimizer paginationCreate() {
-    return new LogicalPlanOptimizer(Arrays.asList(
-        /*
-         * Phase 1: Transformations that rely on relational algebra equivalence
-         */
-        new MergeFilterAndFilter(),
-        new PushFilterUnderSort(),
-        /*
-         * Phase 2: Transformations that rely on data source push down capability
-         */
-        new PushPageSize(),
         new CreatePagingTableScanBuilder(),
         TableScanPushDown.PUSH_DOWN_FILTER,
         TableScanPushDown.PUSH_DOWN_AGGREGATION,

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

@@ -5,45 +5,66 @@
 
 package org.opensearch.sql.planner.optimizer.rule;
 
-import static org.opensearch.sql.planner.optimizer.pattern.Patterns.table;
-
-import com.facebook.presto.matching.Capture;
 import com.facebook.presto.matching.Captures;
 import com.facebook.presto.matching.Pattern;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.List;
 import lombok.Getter;
 import lombok.experimental.Accessors;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalRelation;
 import org.opensearch.sql.planner.optimizer.Rule;
-import org.opensearch.sql.storage.Table;
-import org.opensearch.sql.storage.read.TableScanBuilder;
 
 /**
  * Rule to create a paged TableScanBuilder in pagination request.
  */
-public class CreatePagingTableScanBuilder implements Rule<LogicalRelation> {
-  /** Capture the table inside matched logical relation operator. */
-  private final Capture<Table> capture;
-
+public class CreatePagingTableScanBuilder implements Rule<LogicalPaginate> {
+  /** Capture the table inside matched logical paginate operator. */
+  private LogicalPlan relationParent = null;
   /** Pattern that matches logical relation operator. */
   @Accessors(fluent = true)
   @Getter
-  private final Pattern<LogicalRelation> pattern;
+  private final Pattern<LogicalPaginate> pattern;
 
   /**
    * Constructor.
    */
   public CreatePagingTableScanBuilder() {
-    this.capture = Capture.newCapture();
-    this.pattern = Pattern.typeOf(LogicalRelation.class)
-        .with(table().capturedAs(capture));
+    this.pattern = Pattern.typeOf(LogicalPaginate.class).matching(this::findLogicalRelation);
+  }
+
+  /**
+   * Finds an instance of LogicalRelation and saves a reference in relationParent variable.
+   * @param logicalPaginate An instance of LogicalPaginate
+   * @return true if {@link LogicalRelation} node was found among the descendents of
+   *     {@link this.logicalPaginate}, false otherwise.
+   */
+  private boolean findLogicalRelation(LogicalPaginate logicalPaginate) {
+    Deque<LogicalPlan> plans = new ArrayDeque<>();
+    plans.push(logicalPaginate);
+    do {
+      var plan = plans.pop();
+      if (plan.getChild().stream().anyMatch(LogicalRelation.class::isInstance)) {
+        if (plan.getChild().size() > 1) {
+          throw new UnsupportedOperationException();
+        }
+        relationParent = plan;
+        return true;
+      }
+      plan.getChild().forEach(plans::push);
+    } while (!plans.isEmpty());
+    return false;
   }
 
+
   @Override
-  public LogicalPlan apply(LogicalRelation plan, Captures captures) {
-    TableScanBuilder scanBuilder = captures.get(capture)
-        .createPagedScanBuilder(plan.getPageSize());
-    // TODO: Remove this after Prometheus refactored to new table scan builder too
-    return (scanBuilder == null) ? plan : scanBuilder;
+  public LogicalPlan apply(LogicalPaginate plan, Captures captures) {
+    var logicalRelation = (LogicalRelation) relationParent.getChild().get(0);
+    var scan = logicalRelation.getTable().createPagedScanBuilder(plan.getPageSize());
+    relationParent.replaceChildPlans(List.of(scan));
+
+    return plan;
   }
 }

core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java

@@ -46,9 +46,6 @@ class QueryServiceTest {
   @Mock
   private Planner planner;
 
-  @Mock
-  private Planner paginationPlanner;
-
   @Mock
   private UnresolvedPlan ast;
 
@@ -120,9 +117,8 @@ class Helper {
     public Helper() {
       lenient().when(analyzer.analyze(any(), any())).thenReturn(logicalPlan);
       lenient().when(planner.plan(any())).thenReturn(plan);
-      lenient().when(paginationPlanner.plan(any())).thenReturn(plan);
 
-      queryService = new QueryService(analyzer, executionEngine, planner, paginationPlanner);
+      queryService = new QueryService(analyzer, executionEngine, planner);
     }
 
     Helper executeSuccess() {

core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java

@@ -46,7 +46,7 @@ public static void setUp() {
     when(storageEngine.getTableScan(anyString(), anyString()))
         .thenReturn(mock(TableScanOperator.class));
     paginatedPlanCache = new PaginatedPlanCache(storageEngine);
-    queryService = new QueryService(null, new DefaultExecutionEngine(), null, null);
+    queryService = new QueryService(null, new DefaultExecutionEngine(), null);
   }
 
   @Test

core/src/test/java/org/opensearch/sql/executor/execution/PaginatedPlanTest.java

@@ -26,7 +26,6 @@
 import org.opensearch.sql.executor.QueryService;
 import org.opensearch.sql.planner.Planner;
 import org.opensearch.sql.planner.logical.LogicalPaginate;
-import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
@@ -43,7 +42,7 @@ public static void setUp() {
     when(analyzer.analyze(any(), any())).thenReturn(mock(LogicalPaginate.class));
     var planner = mock(Planner.class);
     when(planner.plan(any())).thenReturn(mock(PhysicalPlan.class));
-    queryService = new QueryService(analyzer, new DefaultExecutionEngine(), null, planner);
+    queryService = new QueryService(analyzer, new DefaultExecutionEngine(), planner);
   }
 
   @Test
@@ -79,7 +78,7 @@ public void onFailure(Exception e) {
       }
     };
     var plan = new PaginatedPlan(QueryId.None, mock(UnresolvedPlan.class), 10,
-        new QueryService(null, new DefaultExecutionEngine(), null, null), listener);
+        new QueryService(null, new DefaultExecutionEngine(), null), listener);
     plan.execute();
   }