Reactivate amortization of plan generation of functions.

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

@@ -23,6 +23,7 @@
 import com.apple.foundationdb.record.query.plan.cascades.RawSqlFunction;
 import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.relational.api.metadata.InvokedRoutine;
+import com.apple.foundationdb.relational.recordlayer.util.MemoizedFunction;
 import com.apple.foundationdb.relational.util.Assert;
 
 import javax.annotation.Nonnull;
@@ -43,7 +44,7 @@ public class RecordLayerInvokedRoutine implements InvokedRoutine {
     private final boolean isTemporary;
 
     @Nonnull
-    private final Function<Boolean, UserDefinedFunction> userDefinedFunctionSupplier;
+    private final Function<Boolean, UserDefinedFunction> userDefinedFunctionProvider;
 
     @Nonnull
     private final UserDefinedFunction serializableFunction;
@@ -52,14 +53,14 @@ public RecordLayerInvokedRoutine(@Nonnull final String description,
                                      @Nonnull final String normalizedDescription,
                                      @Nonnull final String name,
                                      boolean isTemporary,
-                                     @Nonnull final Function<Boolean, UserDefinedFunction> userDefinedFunctionSupplier,
+                                     @Nonnull final Function<Boolean, UserDefinedFunction> userDefinedFunctionProvider,
                                      @Nonnull final UserDefinedFunction serializableFunction) {
         this.description = description;
         this.normalizedDescription = normalizedDescription;
         this.name = name;
         this.isTemporary = isTemporary;
         // TODO this used to be memoized
-        this.userDefinedFunctionSupplier = userDefinedFunctionSupplier;
+        this.userDefinedFunctionProvider = MemoizedFunction.memoize(userDefinedFunctionProvider::apply);
         this.serializableFunction = serializableFunction;
     }
 
@@ -76,8 +77,8 @@ public String getNormalizedDescription() {
     }
 
     @Nonnull
-    public Function<Boolean, UserDefinedFunction> getUserDefinedFunctionSupplier() {
-        return userDefinedFunctionSupplier;
+    public Function<Boolean, UserDefinedFunction> getUserDefinedFunctionProvider() {
+        return userDefinedFunctionProvider;
     }
 
     @Nonnull
@@ -135,15 +136,15 @@ public Builder toBuilder() {
                 .setDescription(getDescription())
                 .setNormalizedDescription(getNormalizedDescription())
                 .setTemporary(isTemporary())
-                .withUserDefinedRoutine(getUserDefinedFunctionSupplier())
+                .withUserDefinedRoutine(getUserDefinedFunctionProvider())
                 .withSerializableFunction(asSerializableFunction());
     }
 
     public static final class Builder {
         private String description;
         private String normalizedDescription;
         private String name;
-        private Function<Boolean, UserDefinedFunction> userDefinedFunctionSupplier;
+        private Function<Boolean, UserDefinedFunction> userDefinedFunctionProvider;
         private UserDefinedFunction serializableFunction;
         private boolean isTemporary;
 
@@ -169,8 +170,8 @@ public Builder setName(@Nonnull final String name) {
         }
 
         @Nonnull
-        public Builder withUserDefinedRoutine(@Nonnull final Function<Boolean, UserDefinedFunction> userDefinedFunctionSupplier) {
-            this.userDefinedFunctionSupplier = userDefinedFunctionSupplier;
+        public Builder withUserDefinedRoutine(@Nonnull final Function<Boolean, UserDefinedFunction> userDefinedFunctionProvider) {
+            this.userDefinedFunctionProvider = userDefinedFunctionProvider;
             return this;
         }
 
@@ -190,10 +191,10 @@ public Builder setTemporary(boolean isTemporary) {
         public RecordLayerInvokedRoutine build() {
             Assert.notNullUnchecked(name);
             Assert.notNullUnchecked(description);
-            Assert.notNullUnchecked(userDefinedFunctionSupplier);
+            Assert.notNullUnchecked(userDefinedFunctionProvider);
             Assert.notNullUnchecked(serializableFunction);
             return new RecordLayerInvokedRoutine(description, normalizedDescription, name, isTemporary,
-                    userDefinedFunctionSupplier, serializableFunction);
+                    userDefinedFunctionProvider, serializableFunction);
         }
     }
 }

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

@@ -629,7 +629,7 @@ public static NormalizationResult normalizeAst(@Nonnull final SchemaTemplate sch
             // immediate materialization of temporary function, this is required to collect any auxiliary literals discovered
             // during plan generation of the temporary function. The literals and combined with query literals and provided
             // for the execution of a (cached) physical plan.
-            final var compiledFunction = (CompiledSqlFunction)recordLayerRoutine.getUserDefinedFunctionSupplier().apply(caseSensitive);
+            final var compiledFunction = (CompiledSqlFunction)recordLayerRoutine.getUserDefinedFunctionProvider().apply(caseSensitive);
             astNormalizer.queryHasherContextBuilder.getLiteralsBuilder().importLiterals(compiledFunction.getAuxiliaryLiterals());
         }
         return new NormalizationResult(

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

@@ -876,12 +876,13 @@ public LogicalOperator resolveTableFunction(@Nonnull final Identifier functionNa
                 : tableFunction.encapsulate(valueArgs);
         if (resultingValue instanceof StreamingValue) {
             final var tableFunctionExpression = new TableFunctionExpression(Assert.castUnchecked(resultingValue, StreamingValue.class));
-            final var resultingQuantifier = Quantifier.forEach(Reference.initialOf(tableFunctionExpression));
+            final var reference = Reference.initialOf(tableFunctionExpression);
+            final var resultingQuantifier = Quantifier.forEach(reference);
             final var output = Expressions.of(LogicalOperator.convertToExpressions(resultingQuantifier));
             return LogicalOperator.newNamedOperator(functionName, output, resultingQuantifier);
         }
-        final var relationalExpression = Assert.castUnchecked(resultingValue, RelationalExpression.class);
-        final var topQun = Quantifier.forEach(Reference.initialOf(relationalExpression));
+        final var tableExpression = Assert.castUnchecked(resultingValue, RelationalExpression.class);
+        final var topQun = Quantifier.forEach(Reference.initialOf(tableExpression));
         return LogicalOperator.newNamedOperator(functionName, Expressions.fromQuantifier(topQun), topQun);
     }
 

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

@@ -25,8 +25,11 @@
 import com.apple.foundationdb.record.query.plan.cascades.Column;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
 import com.apple.foundationdb.record.query.plan.cascades.GraphExpansion;
+import com.apple.foundationdb.record.query.plan.cascades.Memoizer;
+import com.apple.foundationdb.record.query.plan.cascades.PlannerStage;
 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.References;
 import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.record.query.plan.cascades.expressions.RelationalExpression;
 import com.apple.foundationdb.record.query.plan.cascades.expressions.TableFunctionExpression;
@@ -38,12 +41,14 @@
 import com.apple.foundationdb.record.query.plan.cascades.values.StreamingValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.ThrowsValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.Value;
+import com.apple.foundationdb.record.query.plan.cascades.values.translation.ToUniqueAliasesTranslationMap;
 import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
 import com.apple.foundationdb.relational.recordlayer.query.Expression;
 import com.apple.foundationdb.relational.recordlayer.query.Expressions;
 import com.apple.foundationdb.relational.recordlayer.query.Literals;
 import com.apple.foundationdb.relational.util.Assert;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 
 import javax.annotation.Nonnull;
@@ -119,14 +124,8 @@ public RelationalExpression encapsulate(@Nonnull final List<? extends Typed> arg
             }
             resultBuilder.addResultColumn(Column.of(Optional.of(getParameterName(paramIdx)), argumentValue));
         }
-        final var qun = Quantifier.forEach(Reference.initialOf(resultBuilder.addQuantifier(rangeOfOnePlan()).build().buildSelect()),
-                parametersCorrelation.get());
-        final var bodyQun = Quantifier.forEach(Reference.initialOf(body));
-        final var selectBuilder = GraphExpansion.builder()
-                .addQuantifier(bodyQun)
-                .addQuantifier(qun);
-        bodyQun.computeFlowedColumns().forEach(selectBuilder::addResultColumn);
-        return selectBuilder.build().buildSelect();
+        final var argumentsExpression = resultBuilder.addQuantifier(rangeOfOnePlan()).build().buildSelect();
+        return constructTableFunctionExpression(argumentsExpression);
     }
 
     @Nonnull
@@ -155,12 +154,27 @@ public RelationalExpression encapsulate(@Nonnull final Map<String, ? extends Typ
             final var maybePromotedArgument = PromoteValue.inject(argumentValue, computeParameterType(name));
             resultBuilder.addResultColumn(Column.of(Optional.of(name), maybePromotedArgument));
         }
-        final var qun = Quantifier.forEach(Reference.initialOf(resultBuilder.addQuantifier(rangeOfOnePlan()).build().buildSelect()),
-                parametersCorrelation.get());
-        final var bodyQun = Quantifier.forEach(Reference.initialOf(body));
+        final var argumentsExpression = resultBuilder.addQuantifier(rangeOfOnePlan()).build().buildSelect();
+        return constructTableFunctionExpression(argumentsExpression);
+    }
+
+    @Nonnull
+    private RelationalExpression constructTableFunctionExpression(@Nonnull final RelationalExpression argumentsExpression) {
+        final var aliasMap = new ToUniqueAliasesTranslationMap();
+
+        final var bodyRef = Reference.initialOf(body);
+        final var translatedBodyRef = Iterables.getOnlyElement(References.rebaseGraphs(List.of(bodyRef),
+                Memoizer.noMemoization(PlannerStage.INITIAL), aliasMap, false));
+        final var bodyQun = Quantifier.forEach(translatedBodyRef);
+
+        final var translatedParamCorrelation = aliasMap.getSnapshotAliasMap()
+                .getTargetOrDefault(Assert.optionalUnchecked(parametersCorrelation),
+                        Quantifier.uniqueId());
+        final var parametersQun = Quantifier.forEach(Reference.initialOf(argumentsExpression), translatedParamCorrelation);
+
         final var selectBuilder = GraphExpansion.builder()
                 .addQuantifier(bodyQun)
-                .addQuantifier(qun);
+                .addQuantifier(parametersQun);
         bodyQun.computeFlowedColumns().forEach(selectBuilder::addResultColumn);
         return selectBuilder.build().buildSelect();
     }

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

@@ -160,7 +160,7 @@ public static SqlFunctionCatalogImpl newInstance(@Nonnull final RecordLayerSchem
         metadata.getInvokedRoutines().forEach(func ->
                 functionCatalog.registerUserDefinedFunction(
                         Assert.notNullUnchecked(func.getName()),
-                        func.getUserDefinedFunctionSupplier()));
+                        func.getUserDefinedFunctionProvider()));
         return functionCatalog;
     }
 }