Address planner bugs encountered while integrating rewrite rules (#3435)

This contains a number of planner fixes that were first identified when
trying to integrate the planner rewrite rules (see: #3401). They are:

1. There were some values being leaked into the `Traversal` which would
end up polluting the memo structure with additional values that weren't
in the final graph. This adds a sanity check to validate the traversal
contains precisely the information in the tree. There were a bunch of
places where things got out of line, mainly due to something being added
to the memo structure during rule execution (e.g., a new child) but that
was never actually integrated into the final expression (e.g., because
the expression was a duplicate of another in the reference). This adds a
new "garbage collection" like mechanic, where new references are tracked
during rule execution, and then after we're done, any that are still not
in the memo are pruned.
2. Another memoizer issue: The first was that expressions could be
re-used when they were in the wrong planning stage. This would be bad,
and so the memoizer now looks for only references that match the
expected stage.
3. Additionally, we could end up memoizing a reference that had an
incorrect correlation set. Say, for example, that one expression was
added that dropped a correlation (maybe due to predicate
simplification). If we tried to re-use this reference in a place where
that correlation wasn't guaranteed to be defined, which means we'd have
an illegal graph. This addresses this by only looking for pre-existing
references in the memoizer that have the right correlation set. This
also allows us to update `DecorrelateValues` so that it can pick up any
_expression_ that is not sideways correlated without having to worry
about those correlations in any quantifier it pushes down
4. The `PredicateMultiMap` map hit an issue if one tried to add the same
(semantically) equivalent predicate multiple times. This can happen if
the user had the same predicate multiple times (like `a > 4 AND a > 4`)
though it's actually harder to construct than one might think. The
rewrite rules hit it because there was a query like `SELECT * FROM
(SELECT * FROM T WHERE T.a > 4) X WHERE X.a > 4` and then select merge
would be two copies of the same predicate at the same level. The builder
for the `PredicateMultiMap` was always based on `LinkedIdenityMap`, so
it used pointer equality. This updates the built `PredicateMultiMap` to
be based on pointer equality as well instead of semantic equality
(through immutable collections)
5. The `Reference` used to have a check to stop any expression that had
a different correlation set from the reference as a whole from being in
the memo. This is wrong, and it would result in way too many members of
the memo. That check has been removed
6. Expression partitions can now match for the argmin of a tuple of
properties. This allows us to pick the best child select with
tiebreakers to favor expressions that are simpler (with tiebreakers)
7. There was a bug in the expression count property where it would
always report the number of selects on references rather than the number
of the expression requested. This has been updated to correctly return
the expression type desired
8. The `InComparisonToExplodeRule` would allow us to do things like take
an expression like `FROM T WHERE (T.a, T.b) IN ?list` and then explode
the list and turn it into something like `FROM T, ?list x WHERE T.a =
x._0 AND T.b = x._1`. That actually required that the original
comparison had child values, and it wouldn't work if `(T.a, T.b)` could
be re-expressed as some kind of simpler thing. This has been modified to
work on more generic items of type `Record`
9. The exploration rules now only match against exploratory examples
10. Select merge now checks its children to see if pulling them up
introduces duplicate quantifiers, and it uses the new `rebaseGraphs`
method to rename any that are affected if it would do so
11. Constant object values now can be ignored during simplification is
not in the evaluation context. This was mainly for tests, but it would
be fine for actual code as well

As one can see, it's a lot of things. Most of them are actually fairly
minor code changes, though some of them are a bit conceptually tricky

---------

Co-authored-by: normen662 <[email protected]>

Author: alecgrieser

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/EvaluationContext.java

@@ -21,7 +21,6 @@
 package com.apple.foundationdb.record;
 
 import com.apple.foundationdb.annotation.API;
-import com.apple.foundationdb.record.logging.LogMessageKeys;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
 import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
 
@@ -136,24 +135,63 @@ public Object getBinding(@Nonnull final Bindings.Internal type, @Nonnull final C
         return bindings.get(type.bindingName(alias.getId()));
     }
 
+    /**
+     * Whether a value is bound to the provided parameter.
+     *
+     * @param name the name of the parameter to check
+     * @return whether a value is bound to the given parameter
+     * @see Bindings#containsBinding(String) 
+     */
+    public boolean containsBinding(@Nonnull final String name) {
+        return bindings.containsBinding(name);
+    }
+
+    /**
+     * Whether the bindings contain a given special internal correlation.
+     *
+     * @param type the type of the parameter
+     * @param alias the parameter's alias
+     * @return whether the bindings contain that special correlation value
+     */
+    public boolean containsBinding(@Nonnull final Bindings.Internal type, @Nonnull final CorrelationIdentifier alias) {
+        return containsBinding(type.bindingName(alias.getId()));
+    }
+
+    /**
+     * Whether a value is bound to the given constant.
+     *
+     * @param alias the alias of the constant map
+     * @param constantId the identity of the constant within the map
+     * @return whether a value is bound to the given constant
+     */
+    public boolean containsConstantBinding(@Nonnull final CorrelationIdentifier alias, @Nonnull final String constantId) {
+        if (!containsBinding(Bindings.Internal.CONSTANT, alias)) {
+            return false;
+        }
+        final var constantsMap = getConstantsMap(alias);
+        return constantsMap.containsKey(constantId);
+    }
+
+
     /**
      * Dereferences the constant.
      *
      * @param alias the correlation identifier
      * @param constantId the constant id
      * @return de-referenced constant
      */
-    @SuppressWarnings("unchecked")
     @Nullable
     public Object dereferenceConstant(@Nonnull final CorrelationIdentifier alias, @Nonnull final String constantId) {
-        final var constantsMap = (Map<String, ?>)bindings.get(Bindings.Internal.CONSTANT.bindingName(alias.getId()));
-        if (constantsMap == null) {
-            throw new RecordCoreException("could not find constant in the evaluation context")
-                    .addLogInfo(LogMessageKeys.KEY, "'" + alias.getId() + "' - '" + constantId + "'");
-        }
+        final var constantsMap = getConstantsMap(alias);
         return constantsMap.get(constantId);
     }
 
+    @SuppressWarnings("unchecked")
+    @Nonnull
+    private Map<String, ?> getConstantsMap(@Nonnull final CorrelationIdentifier alias) {
+        return (Map<String, ?>) getBinding(Bindings.Internal.CONSTANT, alias);
+    }
+
     @Nonnull
     public TypeRepository getTypeRepository() {
         return typeRepository;

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

@@ -437,6 +437,11 @@ private void planPartial(@Nonnull final Supplier<Reference> referenceSupplier,
                     Debugger.withDebugger(debugger -> debugger.onEvent(nextTask.toTaskEvent(Location.BEGIN)));
                     try {
                         nextTask.execute();
+                        Debugger.sanityCheck(() ->
+                                // This is a fairly expensive check, but it makes sure that the traversal
+                                // is up to date with the query DAG. This is used during memoization, so
+                                // it's important that it has an accurate view of the state of the world.
+                                traversal.verifyIntegrity());
                     } finally {
                         Debugger.withDebugger(debugger -> debugger.onEvent(nextTask.toTaskEvent(Location.END)));
                     }
@@ -1004,6 +1009,12 @@ public void execute() {
         protected void executeRuleCall(@Nonnull CascadesRuleCall ruleCall) {
             ruleCall.run();
 
+            //
+            // During rule construction, references can be added to the memoizer that don't end up in
+            // any final expression. Prune them here
+            //
+            ruleCall.pruneUnusedNewReferences();
+
             //
             // Handle produced artifacts (through yield...() calls)
             //

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

@@ -97,6 +97,8 @@ public class CascadesRuleCall implements ExplorationCascadesRuleCall, Implementa
     @Nonnull
     private final LinkedIdentitySet<PartialMatch> newPartialMatches;
     @Nonnull
+    private final LinkedIdentitySet<Reference> newReferences;
+    @Nonnull
     private final Set<Reference> referencesWithPushedConstraints;
     @Nonnull
     private final EvaluationContext evaluationContext;
@@ -119,6 +121,7 @@ public CascadesRuleCall(@Nonnull final PlannerPhase plannerPhase,
         this.newExploratoryExpressions = new LinkedIdentitySet<>();
         this.newFinalExpressions = new LinkedIdentitySet<>();
         this.newPartialMatches = new LinkedIdentitySet<>();
+        this.newReferences = new LinkedIdentitySet<>();
         this.referencesWithPushedConstraints = Sets.newLinkedHashSet();
         this.evaluationContext = evaluationContext;
     }
@@ -289,6 +292,24 @@ public EvaluationContext getEvaluationContext() {
         return evaluationContext;
     }
 
+    @Nonnull
+    private Reference addNewReference(@Nonnull final Reference newRef) {
+        for (RelationalExpression expression : newRef.getAllMemberExpressions()) {
+            Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expression)));
+            traversal.addExpression(newRef, expression);
+        }
+        newReferences.add(newRef);
+        return newRef;
+    }
+
+    public void pruneUnusedNewReferences() {
+        // Not all of the references created during the course of a rule call end up in a yielded expression.
+        // At the end of rule call execution, remove any of the newly created references that aren't referenced
+        // by another reference in the graph
+        traversal.pruneUnreferencedRefs(newReferences);
+        newReferences.clear();
+    }
+
     @Nonnull
     @Override
     public Reference memoizeExpressions(@Nonnull final Collection<? extends RelationalExpression> exploratoryExpressions,
@@ -335,7 +356,9 @@ public Reference memoizeExploratoryExpressions(@Nonnull final Collection<? exten
         try {
             Preconditions.checkArgument(expressions.stream().noneMatch(expression -> expression instanceof RecordQueryPlan));
 
-            // Pick a candidate expression from the expressions collection. This will be used to do the
+            final Set<CorrelationIdentifier> requiredCorrelations = correlatedTo(expressions);
+
+            // Pick a candidate expression from the expressions collection. This will be used to do the topological check
             final RelationalExpression expression = Iterables.getFirst(expressions, null);
             Verify.verify(expression != null, "should not get null from first element of non-empty expressions collection");
 
@@ -350,6 +373,7 @@ public Reference memoizeExploratoryExpressions(@Nonnull final Collection<? exten
             final var expressionToReferenceMap = new LinkedIdentityMap<RelationalExpression, Reference>();
             referencePathsList.stream()
                     .flatMap(Collection::stream)
+                    .filter(referencePath -> isEligibleForReuse(requiredCorrelations, referencePath))
                     .forEach(referencePath -> {
                         final var referencingExpression = referencePath.getExpression();
                         if (expressionToReferenceMap.containsKey(referencingExpression)) {
@@ -365,6 +389,7 @@ public Reference memoizeExploratoryExpressions(@Nonnull final Collection<? exten
             final List<Set<RelationalExpression>> referencingExpressions = referencePathsList.stream()
                     .map(referencePaths ->
                             referencePaths.stream()
+                                    .filter(referencePath -> isEligibleForReuse(requiredCorrelations, referencePath))
                                     .map(Traversal.ReferencePath::getExpression)
                                     .collect(LinkedIdentitySet.toLinkedIdentitySet()))
                     .collect(ImmutableList.toImmutableList());
@@ -394,27 +419,35 @@ public Reference memoizeExploratoryExpressions(@Nonnull final Collection<? exten
             }
 
             // 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;
+            return addNewReference(Reference.ofExploratoryExpressions(plannerPhase.getTargetPlannerStage(), expressions));
         } finally {
             Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.end()));
         }
     }
 
+    // Only look for references that are: (1) in the right planner stage and (2) have appropriate correlations.
+    // It's important that there aren't any extra correlations as if we choose to re-use a reference that
+    // contains a correlation that the original set of expressions don't have, then we might try to use it
+    // in a place that we are not allowed to
+    private boolean isEligibleForReuse(@Nonnull Set<CorrelationIdentifier> requiredCorrelations,
+                                       @Nonnull Traversal.ReferencePath path) {
+        return path.getReference().getPlannerStage() == plannerPhase.getTargetPlannerStage() && path.getReference().getCorrelatedTo().equals(requiredCorrelations);
+    }
+
     @Nonnull
     private Reference memoizeLeafExpressions(@Nonnull final Collection<? extends RelationalExpression> expressions) {
         Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.begin()));
         try {
             Preconditions.checkArgument(expressions.stream()
                     .allMatch(expression -> !(expression instanceof RecordQueryPlan) && expression.getQuantifiers().isEmpty()));
 
+            final Set<CorrelationIdentifier> requiredCorrelations = correlatedTo(expressions);
             final var leafRefs = traversal.getLeafReferences();
 
             for (final var leafRef : leafRefs) {
+                if (leafRef.getPlannerStage() != plannerPhase.getTargetPlannerStage() || !requiredCorrelations.equals(leafRef.getCorrelatedTo())) {
+                    continue;
+                }
                 if (leafRef.containsAllInMemo(expressions, AliasMap.emptyMap(), false)) {
                     for (RelationalExpression expression : expressions) {
                         Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.reusedExp(expression)));
@@ -423,17 +456,30 @@ private Reference memoizeLeafExpressions(@Nonnull final Collection<? extends Rel
                 }
             }
 
-            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;
+            return addNewReference(Reference.ofExploratoryExpressions(plannerPhase.getTargetPlannerStage(), expressions));
         } finally {
             Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.end()));
         }
     }
 
+    @Nonnull
+    private Set<CorrelationIdentifier> correlatedTo(@Nonnull Collection<? extends RelationalExpression> expressions) {
+        if (expressions.isEmpty()) {
+            return ImmutableSet.of();
+        } else if (expressions.size() == 1) {
+            return Iterables.getOnlyElement(expressions).getCorrelatedTo();
+        } else {
+            // Note: this makes a copy, whereas an approach based on using Sets.union wouldn't need to.
+            // However, we expect there to be a lot of duplicate values amongst the different expressions
+            // here, and we don't want to have to re-do the de-duplication checks like we would with
+            // an approach that uses Sets.union. So we just make a copy here
+            return expressions.stream()
+                    .map(Correlated::getCorrelatedTo)
+                    .flatMap(Collection::stream)
+                    .collect(ImmutableSet.toImmutableSet());
+        }
+    }
+
     @Nonnull
     @Override
     public Reference memoizeFinalExpressionsFromOther(@Nonnull final Reference reference,
@@ -493,12 +539,7 @@ private Reference memoizeExpressionsExactly(@Nonnull final Collection<? extends
         try {
             final var exploratoryExpressionSet = new LinkedIdentitySet<>(exploratoryExpressions);
             final var finalExpressionSet = new LinkedIdentitySet<>(finalExpressions);
-            final var newRef = referenceCreator.apply(exploratoryExpressionSet, finalExpressionSet);
-            for (final var expression : allExpressions) {
-                Debugger.withDebugger(debugger -> debugger.onEvent(InsertIntoMemoEvent.newExp(expression)));
-                traversal.addExpression(newRef, expression);
-            }
-            return newRef;
+            return addNewReference(referenceCreator.apply(exploratoryExpressionSet, finalExpressionSet));
         } finally {
             Debugger.withDebugger(debugger -> allExpressions.forEach(
                     expression -> debugger.onEvent(InsertIntoMemoEvent.end())));

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

@@ -21,7 +21,6 @@
 package com.apple.foundationdb.record.query.plan.cascades;
 
 import com.apple.foundationdb.record.query.plan.cascades.predicates.QueryPredicate;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.SetMultimap;
 
@@ -38,7 +37,7 @@ private PredicateMap(@Nonnull final SetMultimap<QueryPredicate, PredicateMapping
     }
 
     public Optional<PredicateMapping> getMappingOptional(@Nonnull final QueryPredicate queryPredicate) {
-        final ImmutableSet<PredicateMapping> predicateEntries = getMap().get(queryPredicate);
+        final Set<PredicateMapping> predicateEntries = getMap().get(queryPredicate);
         if (predicateEntries.size() != 1) {
             return Optional.empty();
         }

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

@@ -25,7 +25,6 @@
 import com.apple.foundationdb.record.query.plan.cascades.values.Value;
 import com.apple.foundationdb.record.query.plan.cascades.values.translation.PullUp;
 import com.apple.foundationdb.record.query.plan.cascades.values.translation.TranslationMap;
-import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Multimaps;
 import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
@@ -54,7 +53,7 @@ public class PredicateMultiMap {
      * Backing multimap.
      */
     @Nonnull
-    private final ImmutableSetMultimap<QueryPredicate, PredicateMapping> map;
+    private final SetMultimap<QueryPredicate, PredicateMapping> map;
 
     /**
      * Functional interface to reapply a predicate if necessary.
@@ -501,11 +500,13 @@ public PredicateMapping build() {
     }
 
     protected PredicateMultiMap(@Nonnull final SetMultimap<QueryPredicate, PredicateMapping> map) {
-        this.map = ImmutableSetMultimap.copyOf(map);
+        SetMultimap<QueryPredicate, PredicateMapping> copy = Multimaps.newSetMultimap(new LinkedIdentityMap<>(), LinkedIdentitySet::new);
+        map.entries().forEach(entry -> copy.put(entry.getKey(), entry.getValue()));
+        this.map = Multimaps.unmodifiableSetMultimap(copy);
     }
 
     @Nonnull
-    protected ImmutableSetMultimap<QueryPredicate, PredicateMapping> getMap() {
+    protected SetMultimap<QueryPredicate, PredicateMapping> getMap() {
         return map;
     }
 

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

@@ -820,19 +820,6 @@ public Quantifier rebase(@Nonnull final AliasMap translationMap) {
         throw new UnsupportedOperationException("rebase not supported on quantifier");
     }
 
-    @Nonnull
-    @SuppressWarnings("PMD.CompareObjectsWithEquals")
-    public Quantifier translateGraph(@Nonnull final Memoizer memoizer,
-                                     @Nonnull final TranslationMap translationMap,
-                                     final boolean shouldSimplifyValues) {
-        final Reference rangesOver = getRangesOver();
-        final Reference translatedReference =
-                getRangesOver().translateGraph(memoizer, translationMap, shouldSimplifyValues);
-        return rangesOver == translatedReference
-               ? this
-               : overNewReference(translatedReference);
-    }
-
     @Nonnull
     public static CorrelationIdentifier uniqueID() {
         return CorrelationIdentifier.uniqueID(Quantifier.class);