New planner rule added to push predicates down in the query representation (FoundationDB#3324)

This introduces a new query planner rule that will be a part of a larger
framework to canonicalize query representations. The rule pushes down
predicates to ensure that they appear as low in the tree as possible.
So, for example, suppose you had a query like:

```sql
SELECT a, b, c
   FROM (SELECT a, b, c, d FROM T)
   WHERE d = ?
   ORDER BY a, b
```

This will produce a new query expression where the `d = ?` is pushed
lower in the graph, something like:

```sql
SELECT a, b, c
   FROM (SELECT a, b, c, d FROM T where d = ?)
   ORDER BY a, b
```

This simplification helps us during planning to push down predicate
execution to the lowest part of the graph where they can go.

Note that for the moment, this rule has _not_ been added to the default
`PlanningRuleSet`. The reason for this is that some of the rewrite rules
will have weird overlap with some of the other rules, and so we want to
rewrite the query in one phase and then plan and optimize it in another.
The framework for that is not there yet

Author: alecgrieser

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/CascadesRuleCall.java

@@ -33,14 +33,15 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Verify;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 
 import javax.annotation.Nonnull;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Deque;
-import java.util.Objects;
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
@@ -288,22 +289,53 @@ public EvaluationContext getEvaluationContext() {
 
     @Nonnull
     @Override
-    @SuppressWarnings("PMD.CompareObjectsWithEquals")
     public Reference memoizeExploratoryExpression(@Nonnull final RelationalExpression expression) {
-        if (expression.getQuantifiers().isEmpty()) {
-            return memoizeLeafExpression(expression);
+        return memoizeExploratoryExpressions(ImmutableSet.of(expression));
+    }
+
+    @Nonnull
+    @Override
+    @SuppressWarnings("PMD.CompareObjectsWithEquals")
+    public Reference memoizeExploratoryExpressions(@Nonnull final Collection<? extends RelationalExpression> expressions) {
+        Preconditions.checkArgument(!expressions.isEmpty(), "Cannot create reference over empty expression collection");
+        if (expressions.stream().allMatch(expression -> expression.getQuantifiers().isEmpty())) {
+            return memoizeLeafExpressions(expressions);
         }
+        //
+        // This requires that for each expression in the expressions collection, all of its children have already been
+        // placed in the memo structure. We then look to find a
+        //
+        //   1. Pick one variation.
+        //   2. For that variation, find the set of references that point to each of its children. These form a set
+        //      of sets of candidate parents
+        //   3. Take the set intersection of the candidates to get a new candidate set that is made up of all
+        //      expressions that point to _all_ children
+        //   4. For each reference including an expression in the parent candidate set, look to see if that
+        //      candidate already contains all variations in its memo
+        //
+        // In other words, this uses the topology of the graph to find a candidate set of potential references.
+        // Once that has been done, it checks if any of those already contain all variations. Note that we only
+        // need to use one variation for that initial topology check. If we repeated steps 2-4 with the other
+        // variations, then we will either re-tread existing candidate references (which must be missing at
+        // least one variation) or it will be a new reference, but that reference must be missing at least
+        // one child from the first variation and therefore cannot be reused
+        //
         Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.begin()));
         try {
-            Preconditions.checkArgument(!(expression instanceof RecordQueryPlan));
+            Preconditions.checkArgument(expressions.stream().noneMatch(expression -> expression instanceof RecordQueryPlan));
 
-            final var referencePathsList =
-                    expression.getQuantifiers()
-                            .stream()
+            // Pick a candidate expression from the expressions collection. This will be used to do the
+            final RelationalExpression expression = Iterables.getFirst(expressions, null);
+            Verify.verify(expression != null, "should not get null from first element of non-empty expressions collection");
+
+            // For each child quantifier of this expression, look up the set of parent nodes in the
+            // memo structure that point to the (already memoized) child
+            final var referencePathsList = expression.getQuantifiers().stream()
                             .map(Quantifier::getRangesOver)
                             .map(traversal::getParentRefPaths)
                             .collect(ImmutableList.toImmutableList());
 
+            // Gather the references to the expressions that include each child
             final var expressionToReferenceMap = new LinkedIdentityMap<RelationalExpression, Reference>();
             referencePathsList.stream()
                     .flatMap(Collection::stream)
@@ -318,67 +350,73 @@ public Reference memoizeExploratoryExpression(@Nonnull final RelationalExpressio
                         }
                     });
 
-            final var referencingExpressions =
-                    referencePathsList.stream()
-                            .map(referencePaths ->
-                                    referencePaths.stream()
-                                            .map(Traversal.ReferencePath::getExpression)
-                                            .collect(LinkedIdentitySet.toLinkedIdentitySet()))
-                            .collect(ImmutableList.toImmutableList());
+            // From the traversal, get the sets of expressions that point to each child
+            final List<Set<RelationalExpression>> referencingExpressions = referencePathsList.stream()
+                    .map(referencePaths ->
+                            referencePaths.stream()
+                                    .map(Traversal.ReferencePath::getExpression)
+                                    .collect(LinkedIdentitySet.toLinkedIdentitySet()))
+                    .collect(ImmutableList.toImmutableList());
 
+            // Compute the set of expressions that point to _all_ of the needed children
             final var referencingExpressionsIterator = referencingExpressions.iterator();
             final var commonReferencingExpressions = new LinkedIdentitySet<>(referencingExpressionsIterator.next());
             while (referencingExpressionsIterator.hasNext()) {
                 commonReferencingExpressions.retainAll(referencingExpressionsIterator.next());
             }
 
-            commonReferencingExpressions.removeIf(commonReferencingExpression ->
-                    !Reference.isMemoizedExpression(expression, commonReferencingExpression));
-
-            if (!commonReferencingExpressions.isEmpty()) {
-                Debugger.withDebugger(debugger ->
-                        debugger.onEvent(InsertIntoMemoEvent.reusedExpWithReferences(expression,
-                                commonReferencingExpressions.stream()
-                                        .map(expressionToReferenceMap::get)
-                                        .collect(ImmutableList.toImmutableList()))));
-                final var reference =
-                        expressionToReferenceMap.get(Objects.requireNonNull(
-                                Iterables.getFirst(commonReferencingExpressions, null)));
-                Verify.verifyNotNull(reference);
-                Verify.verify(reference != this.root);
-                return reference;
+            // For each reference that has the right topology (that is, it points to all the children),
+            // limit the candidates to those that actually contain all the variations in its memo
+            final List<Reference> existingRefs = commonReferencingExpressions.stream()
+                    .map(expressionToReferenceMap::get)
+                    .filter(ref -> ref.containsAllInMemo(expressions, AliasMap.emptyMap(), false))
+                    .collect(ImmutableList.toImmutableList());
+
+            // If we found such a reference, re-use it
+            if (!existingRefs.isEmpty()) {
+                Reference existingReference = existingRefs.get(0);
+                expressions.forEach(expr ->
+                        Debugger.withDebugger(debugger ->
+                                debugger.onEvent(InsertIntoMemoEvent.reusedExpWithReferences(expr, existingRefs))));
+                Verify.verify(existingReference != this.root);
+                return existingReference;
             }
 
-            Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expression)));
-
-            final var newRef = Reference.ofExploratoryExpression(plannerPhase.getTargetPlannerStage(), expression);
-            traversal.addExpression(newRef, expression);
+            // If we didn't find one, create a new reference and add it to the memo
+            final var newRef = Reference.ofExploratoryExpressions(plannerPhase.getTargetPlannerStage(), expressions);
+            expressions.forEach(expr -> {
+                Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expr)));
+                traversal.addExpression(newRef, expr);
+            });
             return newRef;
         } finally {
             Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.end()));
         }
     }
 
     @Nonnull
-    private Reference memoizeLeafExpression(@Nonnull final RelationalExpression expression) {
+    private Reference memoizeLeafExpressions(@Nonnull final Collection<? extends RelationalExpression> expressions) {
         Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.begin()));
         try {
-            Preconditions.checkArgument(!(expression instanceof RecordQueryPlan));
-            Preconditions.checkArgument(expression.getQuantifiers().isEmpty());
+            Preconditions.checkArgument(expressions.stream()
+                    .allMatch(expression -> !(expression instanceof RecordQueryPlan) && expression.getQuantifiers().isEmpty()));
 
             final var leafRefs = traversal.getLeafReferences();
 
             for (final var leafRef : leafRefs) {
-                for (final var member : leafRef.getExploratoryExpressions()) {
-                    if (Reference.isMemoizedExpression(expression, member)) {
+                if (leafRef.containsAllInMemo(expressions, AliasMap.emptyMap(), false)) {
+                    for (RelationalExpression expression : expressions) {
                         Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.reusedExp(expression)));
-                        return leafRef;
                     }
+                    return leafRef;
                 }
             }
-            Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expression)));
-            final var newRef = Reference.ofExploratoryExpression(plannerPhase.getTargetPlannerStage(), expression);
-            traversal.addExpression(newRef, expression);
+
+            final Reference newRef = Reference.ofExploratoryExpressions(plannerPhase.getTargetPlannerStage(), expressions);
+            for (RelationalExpression expression : expressions) {
+                Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expression)));
+                traversal.addExpression(newRef, expression);
+            }
             return newRef;
         } finally {
             Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.end()));

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/ExploratoryMemoizer.java

@@ -24,6 +24,7 @@
 import com.apple.foundationdb.record.query.plan.cascades.expressions.RelationalExpression;
 
 import javax.annotation.Nonnull;
+import java.util.Collection;
 
 /**
  * Interface to capture the parts of the memoizer API that only exploration rules are allowed to call.
@@ -34,10 +35,12 @@ public interface ExploratoryMemoizer {
      * semantically equivalent to the expression that is passed in, the reference containing the previously memoized
      * expression is returned allowing the planner to reuse that reference. If no such previously memoized expression is
      * found, a new reference is created and returned to the caller.
-     * <br>
+     *
+     * <p>
      * The expression that is passed in must be an exploratory expression of the current planner phase. If this method
      * is used in an incompatible context and a reused reference is returned, the effects of mutating such a reference
      * through other planner logic are undefined and should be avoided.
+     * </p>
      *
      * @param expression exploratory expression to memoize
      *
@@ -46,6 +49,23 @@ public interface ExploratoryMemoizer {
     @Nonnull
     Reference memoizeExploratoryExpression(@Nonnull RelationalExpression expression);
 
+    /**
+     * Memoize the given collection of {@link RelationalExpression}s. If a reference of previously memoized expressions
+     * is found that contains all of the given expressions, then this will return the existing reference. Otherwise,
+     * this will add all of the expressions to the memoization structure and return a new reference.
+     *
+     * <p>
+     * Like with {@link #memoizeExploratoryExpression(RelationalExpression)}, all of the passed in expressions must
+     * be exploratory expressions of the current planner phase. See that method for details.
+     * </p>
+     *
+     * @param expressions the collection of exploratory expressions to memoize
+     * @return a new or reused reference
+     * @see #memoizeExploratoryExpression(RelationalExpression)
+     * */
+    @Nonnull
+    Reference memoizeExploratoryExpressions(@Nonnull Collection<? extends RelationalExpression> expressions);
+
     /**
      * Return a new {@link ReferenceBuilder} for exploratory expressions. The expression passed in is memoized when
      * the builder's {@link ReferenceBuilder#reference()} is called. The expression is treated as an exploratory

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/Reference.java

@@ -421,14 +421,17 @@ boolean containsAllInMemo(@Nonnull final Reference otherRef,
             return true;
         }
 
-        for (final RelationalExpression otherExpression : otherRef.getExploratoryExpressions()) {
-            if (!exploratoryMembers.containsInMemo(otherExpression, equivalenceMap)) {
-                return false;
-            }
-        }
+        return containsAllInMemo(otherRef.getExploratoryExpressions(), equivalenceMap, false)
+                && containsAllInMemo(otherRef.getFinalExpressions(), equivalenceMap, true);
+    }
 
-        for (final RelationalExpression otherExpression : otherRef.getFinalExpressions()) {
-            if (!finalMembers.containsInMemo(otherExpression, equivalenceMap)) {
+    @API(API.Status.INTERNAL)
+    boolean containsAllInMemo(@Nonnull final Collection<? extends RelationalExpression> expressions,
+                              @Nonnull final AliasMap equivalenceMap,
+                              boolean isFinal) {
+        Members members = isFinal ? finalMembers : exploratoryMembers;
+        for (final RelationalExpression otherExpression : expressions) {
+            if (!members.containsInMemo(otherExpression, equivalenceMap)) {
                 return false;
             }
         }
@@ -690,15 +693,6 @@ public String show(final boolean renderSingleGroups) {
         return PlannerGraphVisitor.show(renderSingleGroups, this);
     }
 
-    public static boolean isMemoizedExpression(@Nonnull final RelationalExpression expression,
-                                               @Nonnull final RelationalExpression otherExpression) {
-        if (!expression.getCorrelatedTo().equals(otherExpression.getCorrelatedTo())) {
-            return false;
-        }
-
-        return isMemoizedExpression(expression, otherExpression, AliasMap.emptyMap());
-    }
-
     @SuppressWarnings("PMD.CompareObjectsWithEquals")
     private static boolean isMemoizedExpression(@Nonnull final RelationalExpression member,
                                                 @Nonnull final RelationalExpression otherExpression,
@@ -710,6 +704,14 @@ private static boolean isMemoizedExpression(@Nonnull final RelationalExpression
             return false;
         }
 
+        final Set<CorrelationIdentifier> originalCorrelatedTo = member.getCorrelatedTo();
+        final Set<CorrelationIdentifier> translatedCorrelatedTo = equivalenceMap.definesOnlyIdentities()
+                ? originalCorrelatedTo
+                : originalCorrelatedTo.stream().map(id -> equivalenceMap.getTargetOrDefault(id, id)).collect(ImmutableSet.toImmutableSet());
+        if (!translatedCorrelatedTo.equals(otherExpression.getCorrelatedTo())) {
+            return false;
+        }
+
         final List<? extends Quantifier> quantifiers = member.getQuantifiers();
         final List<? extends Quantifier> otherQuantifiers = otherExpression.getQuantifiers();
         if (member.getQuantifiers().size() != otherQuantifiers.size()) {

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/QuantifierMatchers.java

@@ -21,8 +21,8 @@
 package com.apple.foundationdb.record.query.plan.cascades.matching.structure;
 
 import com.apple.foundationdb.annotation.API;
-import com.apple.foundationdb.record.query.plan.cascades.Reference;
 import com.apple.foundationdb.record.query.plan.cascades.Quantifier;
+import com.apple.foundationdb.record.query.plan.cascades.Reference;
 import com.apple.foundationdb.record.query.plan.cascades.expressions.RelationalExpression;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
 import com.google.common.collect.ImmutableList;
@@ -129,17 +129,29 @@ public static BindingMatcher<Quantifier.ForEach> forEachQuantifierOverRef(@Nonnu
     }
 
     @Nonnull
-    public static BindingMatcher<Quantifier.ForEach> forEachQuantifierWithDefaultOnEmptyOverRef(@Nonnull final BindingMatcher<? extends Reference> downstream) {
+    public static BindingMatcher<Quantifier.ForEach> forEachQuantifierOverRef(@Nonnull final BindingMatcher<? extends Reference> downstream, BindingMatcher<? super Boolean> defaultOnEmptyMatcher) {
         return typedWithDownstream(Quantifier.ForEach.class,
                 Extractor.identity(),
                 AllOfMatcher.matchingAllOf(Quantifier.ForEach.class,
                         ImmutableList.of(
                                 typedWithDownstream(Quantifier.ForEach.class,
                                         Extractor.of(Quantifier.ForEach::isNullOnEmpty, name -> "withDefaultOnEmpty(" + name + ")"),
-                                        PrimitiveMatchers.equalsObject(true)),
+                                        defaultOnEmptyMatcher
+                                ),
                                 typedWithDownstream(Quantifier.ForEach.class,
                                         Extractor.of(Quantifier::getRangesOver, name -> "rangesOver(" + name + ")"),
-                                        downstream))));
+                                        downstream)
+                        )));
+    }
+
+    @Nonnull
+    public static BindingMatcher<Quantifier.ForEach> forEachQuantifierWithDefaultOnEmptyOverRef(@Nonnull final BindingMatcher<? extends Reference> downstream) {
+        return forEachQuantifierOverRef(downstream, PrimitiveMatchers.equalsObject(true));
+    }
+
+    @Nonnull
+    public static BindingMatcher<Quantifier.ForEach> forEachQuantifierWithoutDefaultOnEmptyOverRef(@Nonnull final BindingMatcher<? extends Reference> downstream) {
+        return forEachQuantifierOverRef(downstream, PrimitiveMatchers.equalsObject(false));
     }
 
     @Nonnull

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/RecordQueryPlanMatchers.java

@@ -40,6 +40,7 @@
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryAbstractDataModificationPlan;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryAggregateIndexPlan;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryCoveringIndexPlan;
+import com.apple.foundationdb.record.query.plan.plans.RecordQueryDefaultOnEmptyPlan;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryDeletePlan;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryExplodePlan;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan;
@@ -163,6 +164,11 @@ public static SetMatcher<? extends RecordQueryPlan> exactlyPlansInAnyOrder(@Nonn
         return exactlyInAnyOrder(downstreams);
     }
 
+    @Nonnull
+    public static BindingMatcher<RecordQueryDefaultOnEmptyPlan> defaultOnEmptyPlan(@Nonnull final BindingMatcher<? extends RecordQueryPlan> downstream) {
+        return childrenPlans(RecordQueryDefaultOnEmptyPlan.class, all(downstream));
+    }
+
     @Nonnull
     public static BindingMatcher<RecordQueryFilterPlan> filter(@Nonnull final BindingMatcher<? extends Quantifier> downstream) {
         return ofTypeOwning(RecordQueryFilterPlan.class, any(downstream));