Consolidate with partial pushdown logic after merge

Signed-off-by: Songkan Tang <songkant@amazon.com>

Author: songkant-aws

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -95,6 +95,20 @@ public void supportPartialPushDown_NoPushIfAllFailed() throws IOException {
     assertJsonEqualsIgnoreId(expected, result);
   }
 
+  @Test
+  public void supportPartialPushDownScript() throws IOException {
+    Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled());
+    // field `address` is text type without keyword subfield, so we cannot push it down.
+    // But the second condition can be translated to script, so the second one is pushed down.
+    String query =
+        "source=opensearch-sql_test_index_account | where address = '671 Bristol Street' and age -"
+            + " 2 = 30 | fields firstname, age, address";
+    var result = explainQueryToString(query);
+    String expected =
+        loadFromFile("expectedOutput/calcite/explain_partial_filter_script_push.json");
+    assertJsonEqualsIgnoreId(expected, result);
+  }
+
   // Only for Calcite, as v2 gets unstable serialized string for function
   @Test
   public void testFilterScriptPushDownExplain() throws Exception {

integ-test/src/test/resources/expectedOutput/calcite/explain_partial_filter_script_push.json

@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(firstname=[$1], age=[$8], address=[$2])\n  LogicalFilter(condition=[AND(=($2, '671 Bristol Street'), =(-($8, 2), 30))])\n    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableCalc(expr#0..2=[{inputs}], expr#3=['671 Bristol Street':VARCHAR], expr#4=[=($t1, $t3)], firstname=[$t0], age=[$t2], address=[$t1], $condition=[$t4])\n  CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[firstname, address, age], SCRIPT->=(-($2, 2), 30)], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"query\":{\"bool\":{\"must\":[{\"script\":{\"script\":{\"source\":\"{\\\"langType\\\":\\\"calcite\\\",\\\"script\\\":\\\"rO0ABXNyABFqYXZhLnV0aWwuQ29sbFNlcleOq7Y6G6gRAwABSQADdGFneHAAAAADdwQAAAAGdAAHcm93VHlwZXQBVnsKICAiZmllbGRzIjogWwogICAgewogICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgInByZWNpc2lvbiI6IC0xLAogICAgICAibmFtZSI6ICJmaXJzdG5hbWUiCiAgICB9LAogICAgewogICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgInByZWNpc2lvbiI6IC0xLAogICAgICAibmFtZSI6ICJhZGRyZXNzIgogICAgfSwKICAgIHsKICAgICAgInR5cGUiOiAiQklHSU5UIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgIm5hbWUiOiAiYWdlIgogICAgfQogIF0sCiAgIm51bGxhYmxlIjogZmFsc2UKfXQABGV4cHJ0AnJ7CiAgIm9wIjogewogICAgIm5hbWUiOiAiPSIsCiAgICAia2luZCI6ICJFUVVBTFMiLAogICAgInN5bnRheCI6ICJCSU5BUlkiCiAgfSwKICAib3BlcmFuZHMiOiBbCiAgICB7CiAgICAgICJvcCI6IHsKICAgICAgICAibmFtZSI6ICItIiwKICAgICAgICAia2luZCI6ICJNSU5VUyIsCiAgICAgICAgInN5bnRheCI6ICJCSU5BUlkiCiAgICAgIH0sCiAgICAgICJvcGVyYW5kcyI6IFsKICAgICAgICB7CiAgICAgICAgICAiaW5wdXQiOiAyLAogICAgICAgICAgIm5hbWUiOiAiJDIiCiAgICAgICAgfSwKICAgICAgICB7CiAgICAgICAgICAibGl0ZXJhbCI6IDIsCiAgICAgICAgICAidHlwZSI6IHsKICAgICAgICAgICAgInR5cGUiOiAiSU5URUdFUiIsCiAgICAgICAgICAgICJudWxsYWJsZSI6IGZhbHNlCiAgICAgICAgICB9CiAgICAgICAgfQogICAgICBdLAogICAgICAidHlwZSI6IHsKICAgICAgICAidHlwZSI6ICJCSUdJTlQiLAogICAgICAgICJudWxsYWJsZSI6IHRydWUKICAgICAgfQogICAgfSwKICAgIHsKICAgICAgImxpdGVyYWwiOiAzMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiSU5URUdFUiIsCiAgICAgICAgIm51bGxhYmxlIjogZmFsc2UKICAgICAgfQogICAgfQogIF0KfXQACmZpZWxkVHlwZXNzcgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAAA3QACWZpcnN0bmFtZXNyADpvcmcub3BlbnNlYXJjaC5zcWwub3BlbnNlYXJjaC5kYXRhLnR5cGUuT3BlblNlYXJjaFRleHRUeXBlrYOjkwTjMUQCAAFMAAZmaWVsZHN0AA9MamF2YS91dGlsL01hcDt4cgA6b3JnLm9wZW5zZWFyY2guc3FsLm9wZW5zZWFyY2guZGF0YS50eXBlLk9wZW5TZWFyY2hEYXRhVHlwZcJjvMoC+gU1AgADTAAMZXhwckNvcmVUeXBldAArTG9yZy9vcGVuc2VhcmNoL3NxbC9kYXRhL3R5cGUvRXhwckNvcmVUeXBlO0wAC21hcHBpbmdUeXBldABITG9yZy9vcGVuc2VhcmNoL3NxbC9vcGVuc2VhcmNoL2RhdGEvdHlwZS9PcGVuU2VhcmNoRGF0YVR5cGUkTWFwcGluZ1R5cGU7TAAKcHJvcGVydGllc3EAfgALeHB+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AAdVTktOT1dOfnIARm9yZy5vcGVuc2VhcmNoLnNxbC5vcGVuc2VhcmNoLmRhdGEudHlwZS5PcGVuU2VhcmNoRGF0YVR5cGUkTWFwcGluZ1R5cGUAAAAAAAAAABIAAHhxAH4AEXQABFRleHRzcgA8c2hhZGVkLmNvbS5nb29nbGUuY29tbW9uLmNvbGxlY3QuSW1tdXRhYmxlTWFwJFNlcmlhbGl6ZWRGb3JtAAAAAAAAAAACAAJMAARrZXlzdAASTGphdmEvbGFuZy9PYmplY3Q7TAAGdmFsdWVzcQB+ABh4cHVyABNbTGphdmEubGFuZy5PYmplY3Q7kM5YnxBzKWwCAAB4cAAAAAB1cQB+ABoAAAAAc3EAfgAAAAAAA3cEAAAAAnQAB2tleXdvcmRzcQB+AAx+cQB+ABB0AAZTVFJJTkd+cQB+ABR0AAdLZXl3b3JkcQB+ABl4dAAHYWRkcmVzc3NxAH4ACnEAfgAScQB+ABVxAH4AGXNxAH4AAAAAAAN3BAAAAAB4dAADYWdlfnEAfgAQdAAETE9OR3h4\\\"}\",\"lang\":\"opensearch_compounded_script\",\"params\":{\"utcTimestamp\":*}},\"boost\":1.0}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"firstname\",\"address\",\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java

@@ -177,16 +177,20 @@ public static QueryExpression analyzeExpression(
       List<String> schema,
       Map<String, ExprType> fieldTypes,
       RelDataType rowType,
-      RelOptCluster cluster) throws ExpressionNotAnalyzableException {
+      RelOptCluster cluster)
+      throws ExpressionNotAnalyzableException {
     requireNonNull(expression, "expression");
     try {
       // visits expression tree
       QueryExpression queryExpression =
           (QueryExpression) expression.accept(new Visitor(schema, fieldTypes, rowType, cluster));
       return queryExpression;
     } catch (Throwable e) {
-      Throwables.throwIfInstanceOf(e, UnsupportedOperationException.class);
-      throw new ExpressionNotAnalyzableException("Can't convert " + expression, e);
+      if (e instanceof UnsupportedScriptException) {
+        Throwables.throwIfInstanceOf(e, UnsupportedOperationException.class);
+        throw new ExpressionNotAnalyzableException("Can't convert " + expression, e);
+      }
+      return new ScriptQueryExpression(expression, rowType, fieldTypes, cluster);
     }
   }
 
@@ -604,9 +608,6 @@ private static QueryExpression constructQueryExpressionForSearch(
     }
 
     private QueryExpression andOr(RexCall call) {
-      QueryExpression[] expressions = new QueryExpression[call.getOperands().size()];
-      boolean partial = false;
-      int failedCount = 0;
       // For function isEmpty and isBlank, we implement them via expression `isNull or {@function}`,
       // Unlike `OR` in Java, `SHOULD` in DSL will evaluate both branches and lead to NPE.
       if (call.getKind() == SqlKind.OR
@@ -616,40 +617,41 @@ private QueryExpression andOr(RexCall call) {
         throw new UnsupportedScriptException(
             "DSL will evaluate both branches of OR with isNUll, prevent push-down to avoid NPE");
       }
+
+      QueryExpression[] expressions = new QueryExpression[call.getOperands().size()];
+      PredicateAnalyzerException firstError = null;
+      boolean partial = false;
+      int failedCount = 0;
       for (int i = 0; i < call.getOperands().size(); i++) {
+        RexNode operand = call.getOperands().get(i);
         try {
-          Expression expr = call.getOperands().get(i).accept(this);
-          if (expr instanceof NamedFieldExpression) {
-            // nop currently
-          } else {
-            expressions[i] = (QueryExpression) call.getOperands().get(i).accept(this);
-            // Update or simplify the analyzed node list if it is not partial.
-            if (!expressions[i].isPartial())
-              expressions[i].updateAnalyzedNodes(call.getOperands().get(i));
+          Expression expr = tryAnalyzeOperand(operand);
+          if (expr instanceof QueryExpression) {
+            expressions[i] = (QueryExpression) expr;
+            partial |= expressions[i].isPartial();
           }
-          partial |= expressions[i].isPartial();
         } catch (PredicateAnalyzerException e) {
-          try {
-            expressions[i] =
-                new ScriptQueryExpression(call.getOperands().get(i), rowType, fieldTypes, cluster);
-            if (!expressions[i].isPartial())
-              expressions[i].updateAnalyzedNodes(call.getOperands().get(i));
-          } catch (UnsupportedScriptException ex) {
-            if (call.getKind() == SqlKind.OR) throw ex;
-            partial = true;
+          if (firstError == null) {
+            firstError = e;
           }
-        } catch (UnsupportedScriptException e) {
-          if (call.getKind() == SqlKind.OR) throw e;
           partial = true;
           ++failedCount;
           // If we cannot analyze the operand, wrap the RexNode with UnAnalyzableQueryExpression and
           // record them in the array. We will reuse them later.
-          expressions[i] = new UnAnalyzableQueryExpression(call.getOperands().get(i));
+          expressions[i] = new UnAnalyzableQueryExpression(operand);
         }
       }
 
       switch (call.getKind()) {
         case OR:
+          if (partial) {
+            if (firstError != null) {
+              throw firstError;
+            } else {
+              final String message = format(Locale.ROOT, "Unable to handle call: [%s]", call);
+              throw new PredicateAnalyzerException(message);
+            }
+          }
           return CompoundQueryExpression.or(expressions);
         case AND:
           if (failedCount == call.getOperands().size()) {
@@ -664,6 +666,30 @@ private QueryExpression andOr(RexCall call) {
       }
     }
 
+    private Expression tryAnalyzeOperand(RexNode node) {
+      try {
+        Expression expr = node.accept(this);
+        if (expr instanceof NamedFieldExpression) {
+          return expr;
+        }
+        QueryExpression qe = (QueryExpression) expr;
+        if (!qe.isPartial()) {
+          qe.updateAnalyzedNodes(node);
+        }
+        return qe;
+      } catch (PredicateAnalyzerException firstFailed) {
+        try {
+          QueryExpression qe = new ScriptQueryExpression(node, rowType, fieldTypes, cluster);
+          if (!qe.isPartial()) {
+            qe.updateAnalyzedNodes(node);
+          }
+          return qe;
+        } catch (UnsupportedScriptException secondFailed) {
+          throw new PredicateAnalyzerException(secondFailed);
+        }
+      }
+    }
+
     /**
      * Holder class for a pair of expressions. Used to convert {@code 1 = foo} into {@code foo = 1}
      */
@@ -1297,6 +1323,7 @@ private static String timestampValueForPushDown(String value) {
 
   public static class ScriptQueryExpression extends QueryExpression {
     private final String code;
+    private RexNode analyzedNode;
 
     public ScriptQueryExpression(
         RexNode rexNode,
@@ -1331,12 +1358,12 @@ public QueryBuilder builder() {
 
     @Override
     public List<RexNode> getAnalyzedNodes() {
-      return List.of();
+      return List.of(analyzedNode);
     }
 
     @Override
     public void updateAnalyzedNodes(RexNode rexNode) {
-
+      this.analyzedNode = rexNode;
     }
 
     @Override

opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java

@@ -47,7 +47,6 @@
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType;
 import org.opensearch.sql.opensearch.request.PredicateAnalyzer.ExpressionNotAnalyzableException;
-import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.UnsupportedScriptException;
 
 public class PredicateAnalyzerTest {
   final RelDataTypeFactory typeFactory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
@@ -660,11 +659,11 @@ void equals_throwException_TextWithoutKeyword() {
     final RexInputRef field3 =
         builder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 2);
     RexNode call = builder.makeCall(SqlStdOperatorTable.EQUALS, field3, stringLiteral);
-    IllegalArgumentException exception =
+    ExpressionNotAnalyzableException exception =
         assertThrows(
-            IllegalArgumentException.class,
+            ExpressionNotAnalyzableException.class,
             () -> PredicateAnalyzer.analyze(call, schema, fieldTypes));
-    assertEquals("field name is null or empty", exception.getMessage());
+    assertEquals("Can't convert =($2, 'Hi')", exception.getMessage());
   }
 
   @Test
@@ -679,12 +678,10 @@ void isNullOr_throwException() {
             .build();
     // PPL IS_EMPTY is translated to OR(IS_NULL(arg), IS_EMPTY(arg))
     RexNode call = PPLFuncImpTable.INSTANCE.resolve(builder, BuiltinFunctionName.IS_EMPTY, field2);
-    UnsupportedScriptException exception =
+    ExpressionNotAnalyzableException exception =
         assertThrows(
-            UnsupportedScriptException.class,
+            ExpressionNotAnalyzableException.class,
             () -> PredicateAnalyzer.analyzeExpression(call, schema, fieldTypes, rowType, cluster));
-    assertEquals(
-        "DSL will evaluate both branches of OR with isNUll, prevent push-down to avoid NPE",
-        exception.getMessage());
+    assertEquals("Can't convert OR(IS NULL($1), IS EMPTY($1))", exception.getMessage());
   }
 }