Support NOT EXISTS for nested query (#200)

opendistro-for-elasticsearch · Oct 4, 2019 · 407357a · 407357a
1 parent 6ff810f
commit 407357a
Show file tree

Hide file tree

Showing 11 changed files with 221 additions and 42 deletions.
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/domain/Condition.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/domain/Condition.java
@@ -58,6 +58,7 @@ public enum OPEAR {
         TERM,
         IDS_QUERY,
         NESTED_COMPLEX,
+        NOT_EXISTS_NESTED_COMPLEX,
         CHILDREN_COMPLEX,
         SCRIPT,
         NIN_TERMS,
@@ -133,6 +134,7 @@ public enum OPEAR {
             negatives.put(IS, ISN);
             negatives.put(IN, NIN);
             negatives.put(BETWEEN, NBETWEEN);
+            negatives.put(NESTED_COMPLEX, NOT_EXISTS_NESTED_COMPLEX);
         }
 
         static {

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/query/maker/Maker.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/query/maker/Maker.java
@@ -88,7 +88,8 @@ public abstract class Maker {
 
     private static final Set<Condition.OPEAR> NOT_OPEAR_SET = ImmutableSet.of(
             Condition.OPEAR.N, Condition.OPEAR.NIN, Condition.OPEAR.ISN, Condition.OPEAR.NBETWEEN,
-            Condition.OPEAR.NLIKE, Condition.OPEAR.NIN_TERMS, Condition.OPEAR.NTERM
+            Condition.OPEAR.NLIKE, Condition.OPEAR.NIN_TERMS, Condition.OPEAR.NTERM,
+            Condition.OPEAR.NOT_EXISTS_NESTED_COMPLEX
     );
 
     protected Maker(Boolean isQuery) {
@@ -319,6 +320,7 @@ private ToXContent make(Condition cond, String name, Object value) throws SqlPar
                 toXContent = QueryBuilders.idsQuery(type).addIds(ids);
                 break;
             case NESTED_COMPLEX:
+            case NOT_EXISTS_NESTED_COMPLEX:
                 if (value == null || !(value instanceof Where)) {
                     throw new SqlParseException("unsupported nested condition");
                 }

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/rewriter/nestedfield/SQLClause.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/rewriter/nestedfield/SQLClause.java
@@ -22,6 +22,7 @@
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
 import com.alibaba.druid.sql.ast.expr.SQLInSubQueryExpr;
 import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr;
+import com.alibaba.druid.sql.ast.expr.SQLNotExpr;
 import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
 import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem;
 import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlSelectGroupByExpr;
@@ -83,6 +84,8 @@ SQLMethodInvokeExpr replaceByNestedFunction(SQLExpr expr) {
             }
         } else if (parent instanceof MySqlSelectQueryBlock) {
             ((MySqlSelectQueryBlock) parent).setWhere(nestedFunc);
+        } else if (parent instanceof SQLNotExpr) {
+              ((SQLNotExpr) parent).setExpr(nestedFunc);
         } else {
             throw new IllegalStateException("Unsupported place to use nested field under parent: " + parent);
         }

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/rewriter/nestedfield/Where.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/rewriter/nestedfield/Where.java
@@ -17,6 +17,7 @@
 
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
+import com.alibaba.druid.sql.ast.expr.SQLNotExpr;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
 
 /**
@@ -46,7 +47,7 @@ void rewrite(Scope scope) {
                 right().mergeNestedField(scope);
             }
         }
-        mergeIfHaveTagAndIsRootOfWhere(scope);
+        mergeIfHaveTagAndIsRootOfWhereOrNot(scope);
     }
 
     private boolean isLeftChildCondition() {
@@ -64,11 +65,14 @@ private void useAnyChildTag(Scope scope) {
     }
 
     /**
-     * Merge anyway if the root of WHERE clause be reached
+     * Merge anyway if the root of WHERE clause or {@link SQLNotExpr} be reached.
      */
-    private void mergeIfHaveTagAndIsRootOfWhere(Scope scope) {
-        if (!scope.getConditionTag(expr).isEmpty()
-                && expr.getParent() instanceof MySqlSelectQueryBlock) {
+    private void mergeIfHaveTagAndIsRootOfWhereOrNot(Scope scope) {
+        if (scope.getConditionTag(expr).isEmpty()) {
+            return;
+        }
+        if (expr.getParent() instanceof MySqlSelectQueryBlock
+            || expr.getParent() instanceof SQLNotExpr) {
             mergeNestedField(scope);
         }
     }

diff --git a/...n/java/com/amazon/opendistroforelasticsearch/sql/rewriter/parent/SQLExprParentSetter.java b/...n/java/com/amazon/opendistroforelasticsearch/sql/rewriter/parent/SQLExprParentSetter.java
@@ -17,6 +17,7 @@
 
 import com.alibaba.druid.sql.ast.expr.SQLInListExpr;
 import com.alibaba.druid.sql.ast.expr.SQLInSubQueryExpr;
+import com.alibaba.druid.sql.ast.expr.SQLNotExpr;
 import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
 
@@ -42,4 +43,13 @@ public boolean visit(SQLInListExpr expr) {
         expr.getExpr().setParent(expr);
         return true;
     }
+
+    /**
+     * Fix the expr in {@link SQLNotExpr} without parent.
+     */
+    @Override
+    public boolean visit(SQLNotExpr notExpr) {
+        notExpr.getExpr().setParent(notExpr);
+        return true;
+    }
 }
diff --git a/...ain/java/com/amazon/opendistroforelasticsearch/sql/rewriter/subquery/RewriterContext.java b/...ain/java/com/amazon/opendistroforelasticsearch/sql/rewriter/subquery/RewriterContext.java
@@ -34,7 +34,7 @@
  */
 public class RewriterContext {
     private final Deque<SQLTableSource> tableStack = new ArrayDeque<>();
-    private final Deque<SQLBinaryOpExpr> binaryOpStack = new ArrayDeque<>();
+    private final Deque<SQLExpr> conditionStack = new ArrayDeque<>();
     private final List<SQLInSubQueryExpr> sqlInSubQueryExprs = new ArrayList<>();
     private final List<SQLExistsExpr> sqlExistsExprs = new ArrayList<>();
     private final NestedQueryContext nestedQueryDetector = new NestedQueryContext();
@@ -44,11 +44,11 @@ public SQLTableSource popJoin() {
     }
 
     public SQLExpr popWhere() {
-        return binaryOpStack.pop();
+        return conditionStack.pop();
     }
 
-    public void addWhere(SQLBinaryOpExpr expr) {
-        binaryOpStack.push(expr);
+    public void addWhere(SQLExpr expr) {
+        conditionStack.push(expr);
     }
 
     /**

diff --git a/...in/java/com/amazon/opendistroforelasticsearch/sql/rewriter/subquery/SubQueryRewriter.java b/...in/java/com/amazon/opendistroforelasticsearch/sql/rewriter/subquery/SubQueryRewriter.java
@@ -60,8 +60,12 @@ private SQLExpr convertWhere(SQLExpr expr) {
             return ctx.popWhere();
         } else if (expr instanceof SQLBinaryOpExpr) {
             SQLBinaryOpExpr binaryOpExpr = (SQLBinaryOpExpr) expr;
-            binaryOpExpr.setLeft(convertWhere(binaryOpExpr.getLeft()));
-            binaryOpExpr.setRight(convertWhere(binaryOpExpr.getRight()));
+            SQLExpr left = convertWhere(binaryOpExpr.getLeft());
+            left.setParent(binaryOpExpr);
+            binaryOpExpr.setLeft(left);
+            SQLExpr right = convertWhere(binaryOpExpr.getRight());
+            right.setParent(binaryOpExpr);
+            binaryOpExpr.setRight(right);
         }
         return expr;
     }

diff --git a/...mazon/opendistroforelasticsearch/sql/rewriter/subquery/rewriter/NestedExistsRewriter.java b/...mazon/opendistroforelasticsearch/sql/rewriter/subquery/rewriter/NestedExistsRewriter.java
@@ -20,7 +20,7 @@
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
 import com.alibaba.druid.sql.ast.expr.SQLExistsExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
-import com.alibaba.druid.sql.ast.expr.SQLNullExpr;
+import com.alibaba.druid.sql.ast.expr.SQLNotExpr;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource.JoinType;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
@@ -60,32 +60,43 @@ public NestedExistsRewriter(SQLExistsExpr existsExpr, RewriterContext board) {
     }
 
     /**
-     * The from table must be nested field and
-     * The NOT EXISTS is not supported yet.
+     * The from table must be nested field.
      */
     @Override
     public boolean canRewrite() {
-        return ctx.isNestedQuery(from) && !existsExpr.isNot();
+        return ctx.isNestedQuery(from);
     }
 
     @Override
     public void rewrite() {
         ctx.addJoin(from, JoinType.COMMA);
+        ctx.addWhere(rewriteExistsWhere());
+    }
 
-        SQLBinaryOpExpr nullOp = generateNullOp();
+    private SQLExpr rewriteExistsWhere() {
+        SQLBinaryOpExpr translatedWhere;
+        SQLBinaryOpExpr notMissingOp = buildNotMissingOp();
         if (null == where) {
-            ctx.addWhere(nullOp);
+            translatedWhere = notMissingOp;
         } else if (where instanceof SQLBinaryOpExpr) {
-            ctx.addWhere(and(nullOp, (SQLBinaryOpExpr) where));
+            translatedWhere = and(notMissingOp, (SQLBinaryOpExpr) where);
         } else {
             throw new IllegalStateException("unsupported expression in where " + where.getClass());
         }
+
+        if (existsExpr.isNot()) {
+            SQLNotExpr sqlNotExpr = new SQLNotExpr(translatedWhere);
+            translatedWhere.setParent(sqlNotExpr);
+            return sqlNotExpr;
+        } else {
+            return translatedWhere;
+        }
     }
 
-    private SQLBinaryOpExpr generateNullOp() {
+    private SQLBinaryOpExpr buildNotMissingOp() {
         SQLBinaryOpExpr binaryOpExpr = new SQLBinaryOpExpr();
         binaryOpExpr.setLeft(new SQLIdentifierExpr(from.getAlias()));
-        binaryOpExpr.setRight(new SQLNullExpr());
+        binaryOpExpr.setRight(new SQLIdentifierExpr("MISSING"));
         binaryOpExpr.setOperator(SQLBinaryOperator.IsNot);
 
         return binaryOpExpr;

diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/SubqueryIT.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/SubqueryIT.java
@@ -214,14 +214,57 @@ public void nonCorrelatedExistsParentWhere() throws IOException {
     }
 
     @Test
-    public void nonCorrelatedNotExistsUnsupported() throws IOException {
-        exceptionRule.expect(ResponseException.class);
-        exceptionRule.expectMessage("Unsupported subquery");
+    public void nonCorrelatedNotExists() throws IOException {
         String query = String.format(Locale.ROOT,
                                      "SELECT e.name " +
                                      "FROM %s as e " +
                                      "WHERE NOT EXISTS (SELECT * FROM e.projects as p)",
                                      TEST_INDEX_EMPLOYEE_NESTED);
-        executeQuery(query);
+
+        JSONObject response = executeQuery(query);
+        assertThat(
+                response,
+                hitAll(
+                        kvString("/_source/name", is("Susan Smith")),
+                        kvString("/_source/name", is("John Doe"))
+                )
+        );
+    }
+
+    @Test
+    public void nonCorrelatedNotExistsWhere() throws IOException {
+        String query = String.format(Locale.ROOT,
+                                     "SELECT e.name " +
+                                     "FROM %s as e " +
+                                     "WHERE NOT EXISTS (SELECT * FROM e.projects as p WHERE p.name LIKE 'aurora')",
+                                     TEST_INDEX_EMPLOYEE_NESTED);
+
+        JSONObject response = executeQuery(query);
+        assertThat(
+                response,
+                hitAll(
+                        kvString("/_source/name", is("Susan Smith")),
+                        kvString("/_source/name", is("Jane Smith")),
+                        kvString("/_source/name", is("John Doe"))
+                )
+        );
+    }
+
+    @Test
+    public void nonCorrelatedNotExistsParentWhere() throws IOException {
+        String query = String.format(Locale.ROOT,
+                                     "SELECT e.name " +
+                                     "FROM %s as e " +
+                                     "WHERE NOT EXISTS (SELECT * FROM e.projects as p WHERE p.name LIKE 'security') " +
+                                     "AND e.name LIKE 'smith'",
+                                     TEST_INDEX_EMPLOYEE_NESTED);
+
+        JSONObject response = executeQuery(query);
+        assertThat(
+                response,
+                hitAll(
+                        kvString("/_source/name", is("Susan Smith"))
+                )
+        );
     }
 }
diff --git a/...test/java/com/amazon/opendistroforelasticsearch/sql/unittest/NestedFieldRewriterTest.java b/...test/java/com/amazon/opendistroforelasticsearch/sql/unittest/NestedFieldRewriterTest.java
@@ -28,6 +28,7 @@
 import com.alibaba.druid.sql.ast.statement.SQLUnionQuery;
 import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlSelectGroupByExpr;
 import com.amazon.opendistroforelasticsearch.sql.rewriter.nestedfield.NestedFieldRewriter;
+import com.amazon.opendistroforelasticsearch.sql.util.SqlParserUtils;
 import org.junit.Test;
 
 import java.util.List;
@@ -250,16 +251,24 @@ public void subQueryWitSameAlias() {
     @Test
     public void isNotNull() {
         same(
-                query("SELECT e.name FROM employee as e, e.projects as p WHERE p IS NOT NULL"),
-                query("SELECT name FROM employee WHERE nested(projects, 'projects') IS NOT NULL")
+                query("SELECT e.name " +
+                      "FROM employee as e, e.projects as p " +
+                      "WHERE p IS NOT MISSING"),
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE nested(projects, 'projects') IS NOT MISSING")
         );
     }
 
     @Test
     public void isNotNullAndCondition() {
         same(
-                query("SELECT e.name FROM employee as e, e.projects as p WHERE p IS NOT NULL AND p.name LIKE 'security'"),
-                query("SELECT name FROM employee WHERE nested('projects', projects IS NOT NULL AND projects.name LIKE 'security')")
+                query("SELECT e.name " +
+                      "FROM employee as e, e.projects as p " +
+                      "WHERE p IS NOT MISSING AND p.name LIKE 'security'"),
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE nested('projects', projects IS NOT MISSING AND projects.name LIKE 'security')")
         );
     }
 
@@ -271,6 +280,18 @@ public void multiCondition() {
         );
     }
 
+    @Test
+    public void nestedAndParentCondition() {
+        same(
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE nested(projects, 'projects') IS NOT MISSING AND name LIKE 'security'"),
+                query("SELECT e.name " +
+                      "FROM employee e, e.projects p " +
+                      "WHERE p IS NOT MISSING AND e.name LIKE 'security'")
+        );
+    }
+
     @Test
     public void aggWithWhereOnParent() {
         same(
@@ -465,6 +486,53 @@ public void aggInHavingWithWhereOnNestedOrNested() {
         );
     }
 
+    @Test
+    public void notIsNotNull() {
+        same(
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE not (nested(projects, 'projects') IS NOT MISSING)"),
+                query("SELECT e.name " +
+                      "FROM employee as e, e.projects as p " +
+                      "WHERE not (p IS NOT MISSING)")
+        );
+    }
+
+    @Test
+    public void notIsNotNullAndCondition() {
+        same(
+                query("SELECT e.name " +
+                      "FROM employee as e, e.projects as p " +
+                      "WHERE not (p IS NOT MISSING AND p.name LIKE 'security')"),
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE not nested('projects', projects IS NOT MISSING AND projects.name LIKE 'security')")
+        );
+    }
+
+    @Test
+    public void notMultiCondition() {
+        same(
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE not nested('projects', projects.year = 2016 AND projects.name LIKE 'security')"),
+                query("SELECT e.name " +
+                      "FROM employee as e, e.projects as p " +
+                      "WHERE not (p.year = 2016 and p.name LIKE 'security')")
+        );
+    }
+
+    @Test
+    public void notNestedAndParentCondition() {
+        same(
+                query("SELECT name " +
+                      "FROM employee " +
+                      "WHERE (not nested(projects, 'projects') IS NOT MISSING) AND name LIKE 'security'"),
+                query("SELECT e.name " +
+                      "FROM employee e, e.projects p " +
+                      "WHERE not (p IS NOT MISSING) AND e.name LIKE 'security'")
+        );
+    }
 
     private void noImpact(String sql) {
         same(parse(sql), rewrite(parse(sql)));
@@ -585,7 +653,7 @@ private void assertTable(SQLTableSource expect, SQLTableSource actual) {
      * @return     Node parsed out of sql
      */
     private SQLQueryExpr query(String sql) {
-        SQLQueryExpr expr = parse(sql);
+        SQLQueryExpr expr = SqlParserUtils.parse(sql);
         if (sql.contains("nested")) {
             return expr;
         }