Extract method for collecting subqueries

trinodb · Aug 29, 2016 · 6669215 · 6669215
1 parent e000741
commit 6669215
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 27 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java
@@ -71,6 +71,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.ListMultimap;
 import com.google.common.collect.UnmodifiableIterator;
 
@@ -84,7 +85,6 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.sql.planner.ExpressionInterpreter.evaluateConstantExpression;
 import static com.facebook.presto.sql.tree.Join.Type.INNER;
-import static com.facebook.presto.sql.util.AstUtils.nodeContains;
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
 import static com.facebook.presto.util.Types.checkType;
 import static com.google.common.base.Preconditions.checkArgument;
@@ -303,15 +303,13 @@ else if (firstDependencies.stream().allMatch(right.canResolvePredicate()) && sec
                 Optional.empty());
 
         if (node.getType() != INNER) {
-            List<InPredicate> inPredicateSubqueries = analysis.getInPredicateSubqueries(node);
             for (Expression complexExpression : complexJoinExpressions) {
-                for (InPredicate inPredicate :  inPredicateSubqueries) {
-                   if (nodeContains(complexExpression, inPredicate)) {
-                       throw new SemanticException(
-                               NOT_SUPPORTED,
-                               inPredicate,
-                               "IN with subquery predicate in join condition not supported");
-                   }
+                Set<InPredicate> inPredicates = subqueryPlanner.collectInPredicateSubqueries(complexExpression, node);
+                if (!inPredicates.isEmpty()) {
+                    throw new SemanticException(
+                            NOT_SUPPORTED,
+                            Iterables.getLast(inPredicates),
+                            "IN with subquery predicate in join condition not supported");
                 }
             }
 

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/SubqueryPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/SubqueryPlanner.java
@@ -101,27 +101,36 @@ public PlanBuilder handleSubqueries(PlanBuilder builder, Collection<Expression>
 
     public PlanBuilder handleSubqueries(PlanBuilder builder, Expression expression, Node node)
     {
-        builder = appendInPredicateApplyNodes(
-                builder,
-                analysis.getInPredicateSubqueries(node)
-                        .stream()
-                        .filter(inPredicate -> nodeContains(expression, inPredicate.getValueList()))
-                        .collect(toImmutableSet()));
-        builder = appendScalarSubqueryApplyNodes(
-                builder,
-                analysis.getScalarSubqueries(node)
-                        .stream()
-                        .filter(subquery -> nodeContains(expression, subquery))
-                        .collect(toImmutableSet()));
-        builder = appendExistsSubqueryApplyNodes(
-                builder,
-                analysis.getExistsSubqueries(node)
-                        .stream()
-                        .filter(subquery -> nodeContains(expression, subquery))
-                        .collect(toImmutableSet()));
+        builder = appendInPredicateApplyNodes(builder, collectInPredicateSubqueries(expression, node));
+        builder = appendScalarSubqueryApplyNodes(builder, collectScalarSubqueries(expression, node));
+        builder = appendExistsSubqueryApplyNodes(builder, collectExistsSubqueries(expression, node));
         return builder;
     }
 
+    public Set<InPredicate> collectInPredicateSubqueries(Expression expression, Node node)
+    {
+        return analysis.getInPredicateSubqueries(node)
+                .stream()
+                .filter(inPredicate -> nodeContains(expression, inPredicate.getValueList()))
+                .collect(toImmutableSet());
+    }
+
+    public Set<SubqueryExpression> collectScalarSubqueries(Expression expression, Node node)
+    {
+        return analysis.getScalarSubqueries(node)
+                .stream()
+                .filter(subquery -> nodeContains(expression, subquery))
+                .collect(toImmutableSet());
+    }
+
+    public Set<ExistsPredicate> collectExistsSubqueries(Expression expression, Node node)
+    {
+        return analysis.getExistsSubqueries(node)
+                .stream()
+                .filter(subquery -> nodeContains(expression, subquery))
+                .collect(toImmutableSet());
+    }
+
     private PlanBuilder appendInPredicateApplyNodes(PlanBuilder subPlan, Set<InPredicate> inPredicates)
     {
         for (InPredicate inPredicate : inPredicates) {