Avoid breaking correlations when rewriting the logical query graph.

Author: hatyo

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

@@ -25,12 +25,14 @@
 import com.apple.foundationdb.record.query.plan.cascades.Quantifier;
 import com.apple.foundationdb.record.query.plan.cascades.values.LeafValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.Value;
+import com.google.common.base.Verify;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 
 /**
  * Translation map that allows for rebasing of graphs in a way that the resulting rebased graph is only using unique
@@ -43,8 +45,10 @@ public class ToUniqueAliasesTranslationMap implements TranslationMap {
     @Nonnull
     private final Map<CorrelationIdentifier, CorrelationIdentifier> sourceToTargetMap;
 
-    public ToUniqueAliasesTranslationMap() {
+    private ToUniqueAliasesTranslationMap(@Nonnull final AliasMap identities) {
+        Verify.verify(identities.definesOnlyIdentities());
         this.sourceToTargetMap = new LinkedHashMap<>();
+        identities.forEachMapping(sourceToTargetMap::put);
     }
 
     @Nonnull
@@ -87,4 +91,14 @@ private CorrelationIdentifier computeTargetIfAbsent(@Nonnull final CorrelationId
         return sourceToTargetMap.computeIfAbsent(sourceAlias,
                 ignored0 -> Quantifier.uniqueId());
     }
+
+    @Nonnull
+    public static ToUniqueAliasesTranslationMap newInstance() {
+        return new ToUniqueAliasesTranslationMap(AliasMap.emptyMap());
+    }
+
+    @Nonnull
+    public static ToUniqueAliasesTranslationMap newInstance(@Nonnull final Set<CorrelationIdentifier> constantAliases) {
+        return new ToUniqueAliasesTranslationMap(AliasMap.identitiesFor(constantAliases));
+    }
 }

fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/query/plan/cascades/RuleTestHelper.java

@@ -149,7 +149,7 @@ private TestRuleExecution run(RelationalExpression original, EvaluationContext e
         // rewriting and planning
         final var copiedOriginal =
                 Iterables.getOnlyElement(References.rebaseGraphs(ImmutableList.of(Reference.initialOf(original)),
-                        Memoizer.noMemoization(PlannerStage.INITIAL), new ToUniqueAliasesTranslationMap(), false)).get();
+                        Memoizer.noMemoization(PlannerStage.INITIAL), ToUniqueAliasesTranslationMap.newInstance(), false)).get();
         ensureStage(PlannerStage.CANONICAL, copiedOriginal);
         if (rule instanceof ImplementationCascadesRule) {
             //
@@ -245,7 +245,7 @@ public TestRuleExecution assertYields(RelationalExpression original, EvaluationC
             //
             final var copiedExpected =
                     Iterables.getOnlyElement(References.rebaseGraphs(ImmutableList.of(Reference.initialOf(expression)),
-                            Memoizer.noMemoization(PlannerStage.INITIAL), new ToUniqueAliasesTranslationMap(), false)).get();
+                            Memoizer.noMemoization(PlannerStage.INITIAL), ToUniqueAliasesTranslationMap.newInstance(), false)).get();
             ensureStage(PlannerStage.CANONICAL, copiedExpected);
             expectedListBuilder.add(copiedExpected);
         }

fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/query/plan/cascades/TranslateGraphTest.java

@@ -92,7 +92,7 @@ void rebaseGraphTest() {
         final var qun = Quantifier.forEach(Reference.initialOf(graphExpansionBuilder.build().buildSelect()));
         final var expression = LogicalSortExpression.unsorted(qun);
         final var reference = Reference.initialOf(expression);
-        final var translationMap = new ToUniqueAliasesTranslationMap();
+        final var translationMap = ToUniqueAliasesTranslationMap.newInstance();
         final var translatedReferences =
                 References.rebaseGraphs(ImmutableList.of(reference),
                         Memoizer.noMemoization(PlannerStage.INITIAL),
@@ -159,7 +159,7 @@ void rebaseGraphTestWithDiamond() {
         final var qun = Quantifier.forEach(Reference.initialOf(graphExpansionBuilder.build().buildSelect()));
         final var expression = LogicalSortExpression.unsorted(qun);
         final var reference = Reference.initialOf(expression);
-        final var translationMap = new ToUniqueAliasesTranslationMap();
+        final var translationMap = ToUniqueAliasesTranslationMap.newInstance();
         final var translatedReferences =
                 References.rebaseGraphs(ImmutableList.of(reference),
                         Memoizer.noMemoization(PlannerStage.INITIAL),

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerInvokedRoutine.java

@@ -23,8 +23,8 @@
 import com.apple.foundationdb.record.query.plan.cascades.RawSqlFunction;
 import com.apple.foundationdb.relational.api.metadata.InvokedRoutine;
 import com.apple.foundationdb.relational.recordlayer.query.functions.CompiledSqlFunction;
+import com.apple.foundationdb.relational.recordlayer.util.MemoizedFunction;
 import com.apple.foundationdb.relational.util.Assert;
-import com.google.common.base.Suppliers;
 
 import javax.annotation.Nonnull;
 import java.util.Objects;
@@ -55,10 +55,7 @@ public RecordLayerInvokedRoutine(@Nonnull final String description,
         this.normalizedDescription = normalizedDescription;
         this.name = name;
         this.isTemporary = isTemporary;
-        // TODO this used to be memoized
-        final var memoizedTrue = Suppliers.memoize(() -> compilableSqlFunctionSupplier.apply(true));
-        final var memoizedFalse = Suppliers.memoize(() -> compilableSqlFunctionSupplier.apply(false));
-        this.compilableSqlFunctionSupplier = param -> param ? memoizedTrue.get() : memoizedFalse.get();
+        this.compilableSqlFunctionSupplier = MemoizedFunction.memoize(compilableSqlFunctionSupplier::apply);
     }
 
     @Nonnull

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java

@@ -842,19 +842,20 @@ public LogicalOperator resolveTableFunction(@Nonnull final Identifier functionNa
         final var resultingValue = arguments.allNamedArguments()
                 ? tableFunction.encapsulate(arguments.toNamedArgumentInvocation())
                 : tableFunction.encapsulate(valueArgs);
+        final var correlations = ((Correlated<?>)Assert.castUnchecked(resultingValue, Correlated.class)).getCorrelatedTo();
         if (resultingValue instanceof StreamingValue) {
             final var tableFunctionExpression = new TableFunctionExpression(Assert.castUnchecked(resultingValue, StreamingValue.class));
             final var reference = Reference.initialOf(tableFunctionExpression);
             final var translatedReference = Iterables.getOnlyElement(References.rebaseGraphs(List.of(reference),
-                    Memoizer.noMemoization(PlannerStage.INITIAL), new ToUniqueAliasesTranslationMap(), false));
+                    Memoizer.noMemoization(PlannerStage.INITIAL), ToUniqueAliasesTranslationMap.newInstance(correlations), false));
             final var resultingQuantifier = Quantifier.forEach(translatedReference);
             final var output = Expressions.of(LogicalOperator.convertToExpressions(resultingQuantifier));
             return LogicalOperator.newNamedOperator(functionName, output, resultingQuantifier);
         }
         final var relationalExpression = Assert.castUnchecked(resultingValue, RelationalExpression.class);
         final var reference = Reference.initialOf(relationalExpression);
         final var translatedReference = Iterables.getOnlyElement(References.rebaseGraphs(List.of(reference),
-                Memoizer.noMemoization(PlannerStage.INITIAL), new ToUniqueAliasesTranslationMap(), false));
+                Memoizer.noMemoization(PlannerStage.INITIAL), ToUniqueAliasesTranslationMap.newInstance(correlations), false));
         final var topQun = Quantifier.forEach(translatedReference);
         return LogicalOperator.newNamedOperator(functionName, Expressions.fromQuantifier(topQun), topQun);
     }