Disable limit-then-filter pushdown with Calcite

yuancu · yuancu · commit b1066f84d47b · 2025-05-20T15:13:45.000+08:00
Signed-off-by: Yuanchun Shen &lt;yuanchu@amazon.com&gt;
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java
@@ -109,6 +109,36 @@ public void testLimitPushDownExplain() throws Exception {
                 + "| fields ageMinus"));
   }
 
+  @Test
+  public void testLimitWithFilterPushdownExplain() throws Exception {
+    // TODO: Fix limit-then-filter pushdown without Calcite
+    //  Currently both limit-then-filter and filter-then-limit have the
+    //  limit-then-filter effect
+    String expectedFilterThenLimit =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_filter_then_limit_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_filter_then_limit_push.json");
+    assertJsonEqualsIgnoreRelId(
+        expectedFilterThenLimit,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| where age > 30 "
+                + "| head 5 "
+                + "| fields age"));
+
+    String expectedLimitThenFilter =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_then_filter_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_then_filter_push.json");
+    assertJsonEqualsIgnoreRelId(
+        expectedLimitThenFilter,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 5 "
+                + "| where age > 30 "
+                + "| fields age"));
+  }
+
   @Test
   public void testFillNullPushDownExplain() throws Exception {
     String expected = loadFromFile("expectedOutput/ppl/explain_fillnull_push.json");
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_filter_then_limit_push.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_filter_then_limit_push.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(age=[$8])\n  LogicalSort(fetch=[5])\n    LogicalFilter(condition=[>($8, 30)])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[FILTER->>($8, 30), LIMIT->5, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_limit_then_filter_push.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_limit_then_filter_push.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(age=[$8])\n  LogicalFilter(condition=[>($8, 30)])\n    LogicalSort(fetch=[5])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableCalc(expr#0=[{inputs}], expr#1=[30], expr#2=[>($t0, $t1)], age=[$t0], $condition=[$t2])\n  CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[LIMIT->5, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_then_limit_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_then_limit_push.json
@@ -0,0 +1,17 @@
+{
+  "root": {
+    "name": "ProjectOperator",
+    "description": {
+      "fields": "[age]"
+    },
+    "children": [
+      {
+        "name": "OpenSearchIndexScan",
+        "description": {
+          "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"
+        },
+        "children": []
+      }
+    ]
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_limit_then_filter_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_limit_then_filter_push.json
@@ -0,0 +1,17 @@
+{
+  "root": {
+    "name": "ProjectOperator",
+    "description": {
+      "fields": "[age]"
+    },
+    "children": [
+      {
+        "name": "OpenSearchIndexScan",
+        "description": {
+          "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"
+        },
+        "children": []
+      }
+    ]
+  }
+}
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchFilterIndexScanRule.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchFilterIndexScanRule.java
@@ -4,6 +4,7 @@
  */
 package org.opensearch.sql.opensearch.planner.physical;
 
+import java.util.function.Predicate;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelRule;
 import org.apache.calcite.rel.core.Filter;
@@ -55,7 +56,14 @@ public interface Config extends RelRule.Config {
                         .oneInput(
                             b1 ->
                                 b1.operand(CalciteLogicalIndexScan.class)
-                                    .predicate(OpenSearchIndexScanRule::test)
+                                    .predicate(
+                                        // Filter pushdown is skipped if a limit has already been
+                                        // pushed down because the current DSL cannot correctly
+                                        // handle filter pushdown after limit. Both "limit after
+                                        // filter" and "filter after limit" result in the same
+                                        // limit-after-filter DSL.
+                                        Predicate.not(OpenSearchIndexScanRule::isLimitPushed)
+                                            .and(OpenSearchIndexScanRule::test))
                                     .noInputs()));
 
     @Override
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchIndexScanRule.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchIndexScanRule.java
@@ -18,4 +18,8 @@ static boolean test(CalciteLogicalIndexScan scan) {
     final RelOptTable table = scan.getTable();
     return table.unwrap(OpenSearchIndex.class) != null;
   }
+
+  static boolean isLimitPushed(CalciteLogicalIndexScan scan) {
+    return scan.getPushDownContext().isLimitPushed();
+  }
 }
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java
@@ -197,7 +197,11 @@ public void pushDownSort(List<SortBuilder<?>> sortBuilders) {
 
   /** Pushdown size (limit) and from (offset) to DSL request. */
   public void pushDownLimit(Integer limit, Integer offset) {
-    requestedTotalSize = limit;
+    // If there are multiple limit, we take the minimum among them
+    // E.g. for `source=t | head 10 | head 5`, we take 5
+    // This also ensures that the limit won't exceed the initial default value. (set to
+    // Settings.Key.QUERY_SIZE_LIMIT in OpenSearchIndex)
+    requestedTotalSize = Math.min(limit, requestedTotalSize);
     startFrom = offset;
     sourceBuilder.from(offset).size(limit);
   }
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteIndexScan.java
@@ -98,6 +98,7 @@ public double estimateRowCount(RelMetadataQuery mq) {
   public static class PushDownContext extends ArrayDeque<PushDownAction> {
 
     private boolean isAggregatePushed = false;
+    private boolean isLimitPushed = false;
 
     @Override
     public PushDownContext clone() {
@@ -111,6 +112,9 @@ public boolean add(PushDownAction pushDownAction) {
       if (pushDownAction.type == PushDownType.AGGREGATION) {
         isAggregatePushed = true;
       }
+      if (pushDownAction.type == PushDownType.LIMIT) {
+        isLimitPushed = true;
+      }
       return super.add(pushDownAction);
     }
 
@@ -119,6 +123,10 @@ public boolean isAggregatePushed() {
       isAggregatePushed = !isEmpty() && super.peekLast().type == PushDownType.AGGREGATION;
       return isAggregatePushed;
     }
+
+    public boolean isLimitPushed() {
+      return isLimitPushed;
+    }
   }
 
   protected enum PushDownType {

Original file line number	Diff line number	Diff line change
`@@ -18,4 +18,8 @@ static boolean test(CalciteLogicalIndexScan scan) {`
`18`	`18`	`final RelOptTable table = scan.getTable();`
`19`	`19`	`return table.unwrap(OpenSearchIndex.class) != null;`
`20`	`20`	`}`
	`21`	`+`
	`22`	`+ static boolean isLimitPushed(CalciteLogicalIndexScan scan) {`
	`23`	`+ return scan.getPushDownContext().isLimitPushed();`
	`24`	`+ }`
`21`	`25`	`}`
Original file line number	Diff line number	Diff line change
`@@ -98,6 +98,7 @@ public double estimateRowCount(RelMetadataQuery mq) {`
`98`	`98`	`public static class PushDownContext extends ArrayDeque<PushDownAction> {`
`99`	`99`
`100`	`100`	`private boolean isAggregatePushed = false;`
	`101`	`+ private boolean isLimitPushed = false;`
`101`	`102`
`102`	`103`	`@Override`
`103`	`104`	`public PushDownContext clone() {`
`@@ -111,6 +112,9 @@ public boolean add(PushDownAction pushDownAction) {`
`111`	`112`	`if (pushDownAction.type == PushDownType.AGGREGATION) {`
`112`	`113`	`isAggregatePushed = true;`
`113`	`114`	`}`
	`115`	`+ if (pushDownAction.type == PushDownType.LIMIT) {`
	`116`	`+ isLimitPushed = true;`
	`117`	`+ }`
`114`	`118`	`return super.add(pushDownAction);`
`115`	`119`	`}`
`116`	`120`
`@@ -119,6 +123,10 @@ public boolean isAggregatePushed() {`
`119`	`123`	`isAggregatePushed = !isEmpty() && super.peekLast().type == PushDownType.AGGREGATION;`
`120`	`124`	`return isAggregatePushed;`
`121`	`125`	`}`
	`126`	`+`
	`127`	`+ public boolean isLimitPushed() {`
	`128`	`+ return isLimitPushed;`
	`129`	`+ }`
`122`	`130`	`}`
`123`	`131`
`124`	`132`	`protected enum PushDownType {`