Migrate ImplementIntersectAndExceptAsUnion optimizer to rule

presto-main/src/main/java/io/prestosql/sql/planner/PlanOptimizers.java

@@ -46,7 +46,9 @@
 import io.prestosql.sql.planner.iterative.rule.ExtractSpatialJoins;
 import io.prestosql.sql.planner.iterative.rule.GatherAndMergeWindows;
 import io.prestosql.sql.planner.iterative.rule.ImplementBernoulliSampleAsFilter;
+import io.prestosql.sql.planner.iterative.rule.ImplementExceptAsUnion;
 import io.prestosql.sql.planner.iterative.rule.ImplementFilteredAggregations;
+import io.prestosql.sql.planner.iterative.rule.ImplementIntersectAsUnion;
 import io.prestosql.sql.planner.iterative.rule.ImplementLimitWithTies;
 import io.prestosql.sql.planner.iterative.rule.ImplementOffset;
 import io.prestosql.sql.planner.iterative.rule.InlineProjections;
@@ -270,11 +272,11 @@ public PlanOptimizers(
                 statsCalculator,
                 estimatedExchangesCostCalculator,
                 ImmutableSet.<Rule<?>>builder()
-                    .addAll(new SimplifyExpressions(metadata, typeAnalyzer).rules())
-                    .addAll(new UnwrapCastInComparison(metadata, typeAnalyzer).rules())
-                    .addAll(new RemoveDuplicateConditions().rules())
-                    .addAll(new CanonicalizeExpressions().rules())
-                    .build());
+                        .addAll(new SimplifyExpressions(metadata, typeAnalyzer).rules())
+                        .addAll(new UnwrapCastInComparison(metadata, typeAnalyzer).rules())
+                        .addAll(new RemoveDuplicateConditions().rules())
+                        .addAll(new CanonicalizeExpressions().rules())
+                        .build());
 
         PlanOptimizer predicatePushDown = new StatsRecordingPlanOptimizer(optimizerStats, new PredicatePushDown(metadata, typeAnalyzer));
 
@@ -347,7 +349,14 @@ public PlanOptimizers(
                         estimatedExchangesCostCalculator,
                         ImmutableSet.of(new RemoveRedundantIdentityProjections())),
                 new SetFlatteningOptimizer(),
-                new ImplementIntersectAndExceptAsUnion(),
+                new IterativeOptimizer(
+                        ruleStats,
+                        statsCalculator,
+                        estimatedExchangesCostCalculator,
+                        ImmutableList.of(new ImplementIntersectAndExceptAsUnion()),
+                        ImmutableSet.of(
+                                new ImplementIntersectAsUnion(),
+                                new ImplementExceptAsUnion())),
                 new LimitPushDown(), // Run the LimitPushDown after flattening set operators to make it easier to do the set flattening
                 new PruneUnreferencedOutputs(),
                 inlineProjections,

presto-main/src/main/java/io/prestosql/sql/planner/iterative/rule/ImplementExceptAsUnion.java

@@ -0,0 +1,92 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.sql.planner.iterative.rule;
+
+import com.google.common.collect.ImmutableList;
+import io.prestosql.matching.Captures;
+import io.prestosql.matching.Pattern;
+import io.prestosql.sql.planner.iterative.Rule;
+import io.prestosql.sql.planner.plan.Assignments;
+import io.prestosql.sql.planner.plan.ExceptNode;
+import io.prestosql.sql.planner.plan.FilterNode;
+import io.prestosql.sql.planner.plan.ProjectNode;
+import io.prestosql.sql.tree.Expression;
+import io.prestosql.sql.tree.NotExpression;
+
+import java.util.List;
+
+import static com.google.common.collect.Iterables.getFirst;
+import static io.prestosql.sql.ExpressionUtils.and;
+import static io.prestosql.sql.planner.plan.Patterns.except;
+
+/**
+ * Converts EXCEPT queries into UNION ALL..GROUP BY...WHERE
+ * E.g.:
+ * <pre>
+ *     SELECT a FROM foo
+ *     EXCEPT
+ *     SELECT x FROM bar
+ * </pre>
+ * =>
+ * <pre>
+ *     SELECT a
+ *     FROM
+ *     (
+ *         SELECT a,
+ *         COUNT(foo_marker) AS foo_count,
+ *         COUNT(bar_marker) AS bar_count
+ *         FROM
+ *         (
+ *             SELECT a, true as foo_marker, null as bar_marker
+ *             FROM foo
+ *             UNION ALL
+ *             SELECT x, null as foo_marker, true as bar_marker
+ *             FROM bar
+ *         ) T1
+ *     GROUP BY a
+ *     ) T2
+ *     WHERE foo_count >= 1 AND bar_count = 0;
+ * </pre>
+ */
+public class ImplementExceptAsUnion
+        implements Rule<ExceptNode>
+{
+    private static final Pattern<ExceptNode> PATTERN = except();
+
+    @Override
+    public Pattern<ExceptNode> getPattern()
+    {
+        return PATTERN;
+    }
+
+    @Override
+    public Result apply(ExceptNode node, Captures captures, Context context)
+    {
+        SetOperationNodeTranslator translator = new SetOperationNodeTranslator(context.getSymbolAllocator(), context.getIdAllocator());
+        SetOperationNodeTranslator.TranslationResult result = translator.makeSetContainmentPlan(node);
+
+        ImmutableList.Builder<Expression> predicatesBuilder = ImmutableList.builder();
+        List<Expression> presentExpression = result.getPresentExpressions();
+        predicatesBuilder.add(getFirst(presentExpression, null));
+        for (int i = 1; i < presentExpression.size(); i++) {
+            predicatesBuilder.add(new NotExpression(presentExpression.get(i)));
+        }
+
+        return Result.ofPlanNode(
+                new ProjectNode(
+                        context.getIdAllocator().getNextId(),
+                        new FilterNode(context.getIdAllocator().getNextId(), result.getPlanNode(), and(predicatesBuilder.build())),
+                        Assignments.identity(node.getOutputSymbols())));
+    }
+}

presto-main/src/main/java/io/prestosql/sql/planner/iterative/rule/ImplementIntersectAsUnion.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.sql.planner.iterative.rule;
+
+import io.prestosql.matching.Captures;
+import io.prestosql.matching.Pattern;
+import io.prestosql.sql.planner.iterative.Rule;
+import io.prestosql.sql.planner.plan.Assignments;
+import io.prestosql.sql.planner.plan.FilterNode;
+import io.prestosql.sql.planner.plan.IntersectNode;
+import io.prestosql.sql.planner.plan.ProjectNode;
+
+import static io.prestosql.sql.ExpressionUtils.and;
+import static io.prestosql.sql.planner.plan.Patterns.intersect;
+
+/**
+ * Converts INTERSECT queries into UNION ALL..GROUP BY...WHERE
+ * E.g.:
+ * <pre>
+ *     SELECT a FROM foo
+ *     INTERSECT
+ *     SELECT x FROM bar
+ * </pre>
+ * =>
+ * <pre>
+ *     SELECT a
+ *     FROM
+ *     (
+ *         SELECT a,
+ *         COUNT(foo_marker) AS foo_count,
+ *         COUNT(bar_marker) AS bar_count
+ *         FROM
+ *         (
+ *             SELECT a, true as foo_marker, null as bar_marker
+ *             FROM foo
+ *             UNION ALL
+ *             SELECT x, null as foo_marker, true as bar_marker
+ *             FROM bar
+ *         ) T1
+ *     GROUP BY a
+ *     ) T2
+ *     WHERE foo_count >= 1 AND bar_count >= 1;
+ * </pre>
+ */
+public class ImplementIntersectAsUnion
+        implements Rule<IntersectNode>
+{
+    private static final Pattern<IntersectNode> PATTERN = intersect();
+
+    @Override
+    public Pattern<IntersectNode> getPattern()
+    {
+        return PATTERN;
+    }
+
+    @Override
+    public Result apply(IntersectNode node, Captures captures, Context context)
+    {
+        SetOperationNodeTranslator translator = new SetOperationNodeTranslator(context.getSymbolAllocator(), context.getIdAllocator());
+        SetOperationNodeTranslator.TranslationResult result = translator.makeSetContainmentPlan(node);
+
+        return Result.ofPlanNode(
+                new ProjectNode(
+                        context.getIdAllocator().getNextId(),
+                        new FilterNode(context.getIdAllocator().getNextId(), result.getPlanNode(), and(result.getPresentExpressions())),
+                        Assignments.identity(node.getOutputSymbols())));
+    }
+}

presto-main/src/main/java/io/prestosql/sql/planner/iterative/rule/SetOperationNodeTranslator.java

@@ -0,0 +1,183 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.sql.planner.iterative.rule;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableMap;
+import io.prestosql.metadata.Signature;
+import io.prestosql.spi.type.StandardTypes;
+import io.prestosql.spi.type.Type;
+import io.prestosql.sql.planner.PlanNodeIdAllocator;
+import io.prestosql.sql.planner.Symbol;
+import io.prestosql.sql.planner.SymbolAllocator;
+import io.prestosql.sql.planner.plan.AggregationNode;
+import io.prestosql.sql.planner.plan.Assignments;
+import io.prestosql.sql.planner.plan.PlanNode;
+import io.prestosql.sql.planner.plan.ProjectNode;
+import io.prestosql.sql.planner.plan.SetOperationNode;
+import io.prestosql.sql.planner.plan.UnionNode;
+import io.prestosql.sql.tree.Cast;
+import io.prestosql.sql.tree.ComparisonExpression;
+import io.prestosql.sql.tree.Expression;
+import io.prestosql.sql.tree.FunctionCall;
+import io.prestosql.sql.tree.GenericLiteral;
+import io.prestosql.sql.tree.Literal;
+import io.prestosql.sql.tree.NullLiteral;
+import io.prestosql.sql.tree.QualifiedName;
+import io.prestosql.sql.tree.SymbolReference;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.Iterables.concat;
+import static io.prestosql.metadata.FunctionKind.AGGREGATE;
+import static io.prestosql.spi.type.BigintType.BIGINT;
+import static io.prestosql.spi.type.BooleanType.BOOLEAN;
+import static io.prestosql.spi.type.TypeSignature.parseTypeSignature;
+import static io.prestosql.sql.planner.plan.AggregationNode.singleGroupingSet;
+import static io.prestosql.sql.tree.BooleanLiteral.TRUE_LITERAL;
+import static io.prestosql.sql.tree.ComparisonExpression.Operator.GREATER_THAN_OR_EQUAL;
+import static java.util.Objects.requireNonNull;
+
+public class SetOperationNodeTranslator
+{
+    private static final String MARKER = "marker";
+    private static final Signature COUNT_AGGREGATION = new Signature("count", AGGREGATE, parseTypeSignature(StandardTypes.BIGINT), parseTypeSignature(StandardTypes.BOOLEAN));
+    private static final Literal GENERIC_LITERAL = new GenericLiteral("BIGINT", "1");
+    private final SymbolAllocator symbolAllocator;
+    private final PlanNodeIdAllocator idAllocator;
+
+    public SetOperationNodeTranslator(SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator)
+    {
+        this.symbolAllocator = requireNonNull(symbolAllocator, "SymbolAllocator is null");
+        this.idAllocator = requireNonNull(idAllocator, "PlanNodeIdAllocator is null");
+    }
+
+    public TranslationResult makeSetContainmentPlan(SetOperationNode node)
+    {
+        checkArgument(!(node instanceof UnionNode), "Cannot simplify a UnionNode");
+        List<Symbol> markers = allocateSymbols(node.getSources().size(), MARKER, BOOLEAN);
+        // identity projection for all the fields in each of the sources plus marker columns
+        List<PlanNode> withMarkers = appendMarkers(markers, node.getSources(), node);
+
+        // add a union over all the rewritten sources. The outputs of the union have the same name as the
+        // original intersect node
+        List<Symbol> outputs = node.getOutputSymbols();
+        UnionNode union = union(withMarkers, ImmutableList.copyOf(concat(outputs, markers)));
+
+        // add count aggregations and filter rows where any of the counts is >= 1
+        List<Symbol> aggregationOutputs = allocateSymbols(markers.size(), "count", BIGINT);
+        AggregationNode aggregation = computeCounts(union, outputs, markers, aggregationOutputs);
+        List<Expression> presentExpression = aggregationOutputs.stream()
+                .map(symbol -> new ComparisonExpression(GREATER_THAN_OR_EQUAL, symbol.toSymbolReference(), GENERIC_LITERAL))
+                .collect(toImmutableList());
+        return new TranslationResult(aggregation, presentExpression);
+    }
+
+    private List<Symbol> allocateSymbols(int count, String nameHint, Type type)
+    {
+        ImmutableList.Builder<Symbol> symbolsBuilder = ImmutableList.builder();
+        for (int i = 0; i < count; i++) {
+            symbolsBuilder.add(symbolAllocator.newSymbol(nameHint, type));
+        }
+        return symbolsBuilder.build();
+    }
+
+    private List<PlanNode> appendMarkers(List<Symbol> markers, List<PlanNode> nodes, SetOperationNode node)
+    {
+        ImmutableList.Builder<PlanNode> result = ImmutableList.builder();
+        for (int i = 0; i < nodes.size(); i++) {
+            result.add(appendMarkers(idAllocator, symbolAllocator, nodes.get(i), i, markers, node.sourceSymbolMap(i)));
+        }
+        return result.build();
+    }
+
+    private static PlanNode appendMarkers(PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator, PlanNode source, int markerIndex, List<Symbol> markers, Map<Symbol, SymbolReference> projections)
+    {
+        Assignments.Builder assignments = Assignments.builder();
+        // add existing intersect symbols to projection
+        for (Map.Entry<Symbol, SymbolReference> entry : projections.entrySet()) {
+            Symbol symbol = symbolAllocator.newSymbol(entry.getKey().getName(), symbolAllocator.getTypes().get(entry.getKey()));
+            assignments.put(symbol, entry.getValue());
+        }
+
+        // add extra marker fields to the projection
+        for (int i = 0; i < markers.size(); ++i) {
+            Expression expression = (i == markerIndex) ? TRUE_LITERAL : new Cast(new NullLiteral(), StandardTypes.BOOLEAN);
+            assignments.put(symbolAllocator.newSymbol(markers.get(i).getName(), BOOLEAN), expression);
+        }
+
+        return new ProjectNode(idAllocator.getNextId(), source, assignments.build());
+    }
+
+    private UnionNode union(List<PlanNode> nodes, List<Symbol> outputs)
+    {
+        ImmutableListMultimap.Builder<Symbol, Symbol> outputsToInputs = ImmutableListMultimap.builder();
+        for (PlanNode source : nodes) {
+            for (int i = 0; i < source.getOutputSymbols().size(); i++) {
+                outputsToInputs.put(outputs.get(i), source.getOutputSymbols().get(i));
+            }
+        }
+
+        return new UnionNode(idAllocator.getNextId(), nodes, outputsToInputs.build(), outputs);
+    }
+
+    private AggregationNode computeCounts(UnionNode sourceNode, List<Symbol> originalColumns, List<Symbol> markers, List<Symbol> aggregationOutputs)
+    {
+        ImmutableMap.Builder<Symbol, AggregationNode.Aggregation> aggregations = ImmutableMap.builder();
+
+        for (int i = 0; i < markers.size(); i++) {
+            Symbol output = aggregationOutputs.get(i);
+            aggregations.put(output, new AggregationNode.Aggregation(
+                    new FunctionCall(QualifiedName.of("count"), ImmutableList.of(markers.get(i).toSymbolReference())),
+                    COUNT_AGGREGATION,
+                    Optional.empty()));
+        }
+
+        return new AggregationNode(idAllocator.getNextId(),
+                sourceNode,
+                aggregations.build(),
+                singleGroupingSet(originalColumns),
+                ImmutableList.of(),
+                AggregationNode.Step.SINGLE,
+                Optional.empty(),
+                Optional.empty());
+    }
+
+    public static class TranslationResult
+    {
+        private final PlanNode planNode;
+        private final List<Expression> presentExpressions;
+
+        public TranslationResult(PlanNode planNode, List<Expression> presentExpressions)
+        {
+            this.planNode = requireNonNull(planNode, "AggregationNode is null");
+            this.presentExpressions = ImmutableList.copyOf(requireNonNull(presentExpressions, "AggregationOutputs is null"));
+        }
+
+        public PlanNode getPlanNode()
+        {
+            return this.planNode;
+        }
+
+        public List<Expression> getPresentExpressions()
+        {
+            return presentExpressions;
+        }
+    }
+}