Introduce getSqlInvokedFunctions SPI and BuiltInPluginFunctionNamespaceManager for registering sql invoked functions

Author: pdabre12

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionHandle.java

@@ -15,6 +15,7 @@
 
 import com.facebook.presto.common.CatalogSchemaName;
 import com.facebook.presto.common.type.TypeSignature;
+import com.facebook.presto.spi.function.BuiltInFunctionKind;
 import com.facebook.presto.spi.function.FunctionHandle;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.Signature;
@@ -24,18 +25,26 @@
 import java.util.List;
 import java.util.Objects;
 
+import static com.facebook.presto.spi.function.BuiltInFunctionKind.ENGINE;
 import static java.util.Objects.requireNonNull;
 
 public class BuiltInFunctionHandle
         implements FunctionHandle
 {
     private final Signature signature;
+    private final BuiltInFunctionKind builtInFunctionKind;
 
     @JsonCreator
     public BuiltInFunctionHandle(@JsonProperty("signature") Signature signature)
+    {
+        this(signature, ENGINE);
+    }
+
+    public BuiltInFunctionHandle(Signature signature, BuiltInFunctionKind builtInFunctionKind)
     {
         this.signature = requireNonNull(signature, "signature is null");
         checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature);
+        this.builtInFunctionKind = requireNonNull(builtInFunctionKind, "builtInFunctionKind is null");
     }
 
     @JsonProperty
@@ -68,6 +77,12 @@ public CatalogSchemaName getCatalogSchemaName()
         return signature.getName().getCatalogSchemaName();
     }
 
+    @JsonProperty
+    public BuiltInFunctionKind getBuiltInFunctionKind()
+    {
+        return builtInFunctionKind;
+    }
+
     @Override
     public boolean equals(Object o)
     {
@@ -78,13 +93,14 @@ public boolean equals(Object o)
             return false;
         }
         BuiltInFunctionHandle that = (BuiltInFunctionHandle) o;
-        return Objects.equals(signature, that.signature);
+        return Objects.equals(signature, that.signature)
+                && Objects.equals(builtInFunctionKind, that.builtInFunctionKind);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(signature);
+        return Objects.hash(signature, builtInFunctionKind);
     }
 
     @Override

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java

@@ -0,0 +1,259 @@
+/*
+ * 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 com.facebook.presto.metadata;
+
+import com.facebook.presto.common.Page;
+import com.facebook.presto.common.QualifiedObjectName;
+import com.facebook.presto.common.block.BlockEncodingSerde;
+import com.facebook.presto.common.function.SqlFunctionResult;
+import com.facebook.presto.common.type.TypeManager;
+import com.facebook.presto.common.type.UserDefinedType;
+import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.function.AlterRoutineCharacteristics;
+import com.facebook.presto.spi.function.FunctionHandle;
+import com.facebook.presto.spi.function.FunctionMetadata;
+import com.facebook.presto.spi.function.FunctionNamespaceManager;
+import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
+import com.facebook.presto.spi.function.Parameter;
+import com.facebook.presto.spi.function.ScalarFunctionImplementation;
+import com.facebook.presto.spi.function.Signature;
+import com.facebook.presto.spi.function.SqlFunction;
+import com.facebook.presto.spi.function.SqlInvokedFunction;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation;
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.util.concurrent.UncheckedExecutionException;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+
+import static com.facebook.presto.spi.function.BuiltInFunctionKind.PLUGIN;
+import static com.facebook.presto.spi.function.FunctionImplementationType.SQL;
+import static com.facebook.presto.spi.function.FunctionKind.SCALAR;
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Throwables.throwIfInstanceOf;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Collections.emptyList;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.TimeUnit.HOURS;
+
+public class BuiltInPluginFunctionNamespaceManager
+        implements FunctionNamespaceManager<SqlFunction>
+{
+    private volatile FunctionMap functions = new FunctionMap();
+    private final FunctionAndTypeManager functionAndTypeManager;
+    private final Supplier<FunctionMap> cachedFunctions =
+            Suppliers.memoize(this::checkForNamingConflicts);
+    private final LoadingCache<Signature, SpecializedFunctionKey> specializedFunctionKeyCache;
+    private final LoadingCache<SpecializedFunctionKey, ScalarFunctionImplementation> specializedScalarCache;
+
+    public BuiltInPluginFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager)
+    {
+        this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
+        specializedFunctionKeyCache = CacheBuilder.newBuilder()
+                .maximumSize(1000)
+                .expireAfterWrite(1, HOURS)
+                .build(CacheLoader.from(this::doGetSpecializedFunctionKey));
+        specializedScalarCache = CacheBuilder.newBuilder()
+                .maximumSize(1000)
+                .expireAfterWrite(1, HOURS)
+                .build(CacheLoader.from(key -> {
+                    checkArgument(
+                            key.getFunction() instanceof SqlInvokedFunction,
+                            "Unsupported scalar function class: %s",
+                            key.getFunction().getClass());
+                    return new SqlInvokedScalarFunctionImplementation(((SqlInvokedFunction) key.getFunction()).getBody());
+                }));
+    }
+
+    public synchronized void registerPluginFunctions(List<? extends SqlFunction> functions)
+    {
+        checkForNamingConflicts(functions);
+        this.functions = new FunctionMap(this.functions, functions);
+    }
+
+    @Override
+    public FunctionHandle getFunctionHandle(Optional<? extends FunctionNamespaceTransactionHandle> transactionHandle, Signature signature)
+    {
+        return new BuiltInFunctionHandle(signature, PLUGIN);
+    }
+
+    @Override
+    public Collection<SqlFunction> getFunctions(Optional<? extends FunctionNamespaceTransactionHandle> transactionHandle, QualifiedObjectName functionName)
+    {
+        if (functions.list().isEmpty() ||
+                (!functionName.getCatalogSchemaName().equals(functionAndTypeManager.getDefaultNamespace()))) {
+            return emptyList();
+        }
+        return cachedFunctions.get().get(functionName);
+    }
+
+    /**
+     * likePattern / escape is not used for optimization, returning all functions.
+     */
+    @Override
+    public Collection<SqlFunction> listFunctions(Optional<String> likePattern, Optional<String> escape)
+    {
+        return cachedFunctions.get().list();
+    }
+
+    public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle)
+    {
+        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
+        Signature signature = ((BuiltInFunctionHandle) functionHandle).getSignature();
+        SpecializedFunctionKey functionKey;
+        try {
+            functionKey = specializedFunctionKeyCache.getUnchecked(signature);
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+        SqlFunction function = functionKey.getFunction();
+        checkArgument(function instanceof SqlInvokedFunction, "BuiltInPluginFunctionNamespaceManager only support SqlInvokedFunctions");
+        SqlInvokedFunction sqlFunction = (SqlInvokedFunction) function;
+        List<String> argumentNames = sqlFunction.getParameters().stream().map(Parameter::getName).collect(toImmutableList());
+        return new FunctionMetadata(
+                signature.getName(),
+                signature.getArgumentTypes(),
+                argumentNames,
+                signature.getReturnType(),
+                signature.getKind(),
+                sqlFunction.getRoutineCharacteristics().getLanguage(),
+                SQL,
+                function.isDeterministic(),
+                function.isCalledOnNullInput(),
+                sqlFunction.getVersion(),
+                sqlFunction.getComplexTypeFunctionDescriptor());
+    }
+
+    public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHandle functionHandle)
+    {
+        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
+        return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature());
+    }
+
+    @Override
+    public void setBlockEncodingSerde(BlockEncodingSerde blockEncodingSerde)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public FunctionNamespaceTransactionHandle beginTransaction()
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void commit(FunctionNamespaceTransactionHandle transactionHandle)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void abort(FunctionNamespaceTransactionHandle transactionHandle)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void createFunction(SqlInvokedFunction function, boolean replace)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void dropFunction(QualifiedObjectName functionName, Optional parameterTypes, boolean exists)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support drop function");
+    }
+
+    @Override
+    public void alterFunction(QualifiedObjectName functionName, Optional parameterTypes, AlterRoutineCharacteristics alterRoutineCharacteristics)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not alter function");
+    }
+
+    @Override
+    public void addUserDefinedType(UserDefinedType userDefinedType)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support adding user defined types");
+    }
+
+    @Override
+    public Optional<UserDefinedType> getUserDefinedType(QualifiedObjectName typeName)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support getting user defined types");
+    }
+
+    @Override
+    public CompletableFuture<SqlFunctionResult> executeFunction(String source, FunctionHandle functionHandle, Page input, List channels, TypeManager typeManager)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not execute function");
+    }
+
+    private ScalarFunctionImplementation getScalarFunctionImplementation(Signature signature)
+    {
+        checkArgument(signature.getKind() == SCALAR, "%s is not a scalar function", signature);
+        checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature);
+
+        try {
+            return specializedScalarCache.getUnchecked(getSpecializedFunctionKey(signature));
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+    }
+
+    private synchronized FunctionMap checkForNamingConflicts()
+    {
+        Optional<FunctionNamespaceManager<?>> functionNamespaceManager =
+                functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace());
+        checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName());
+        checkForNamingConflicts(functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty()));
+        return functions;
+    }
+
+    private synchronized void checkForNamingConflicts(Collection<? extends SqlFunction> functions)
+    {
+        for (SqlFunction function : functions) {
+            for (SqlFunction existingFunction : this.functions.list()) {
+                checkArgument(!function.getSignature().equals(existingFunction.getSignature()), "Function already registered: %s", function.getSignature());
+            }
+        }
+    }
+
+    private SpecializedFunctionKey doGetSpecializedFunctionKey(Signature signature)
+    {
+        return functionAndTypeManager.getSpecializedFunctionKey(signature, getFunctions(Optional.empty(), signature.getName()));
+    }
+
+    private SpecializedFunctionKey getSpecializedFunctionKey(Signature signature)
+    {
+        try {
+            return specializedFunctionKeyCache.getUnchecked(signature);
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+    }
+}

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -292,9 +292,6 @@
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableListMultimap;
-import com.google.common.collect.Multimap;
-import com.google.common.collect.Multimaps;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import io.airlift.slice.Slice;
 
@@ -1396,44 +1393,6 @@ private static class EmptyTransactionHandle
     {
     }
 
-    private static class FunctionMap
-    {
-        private final Multimap<QualifiedObjectName, SqlFunction> functions;
-
-        public FunctionMap()
-        {
-            functions = ImmutableListMultimap.of();
-        }
-
-        public FunctionMap(FunctionMap map, Iterable<? extends SqlFunction> functions)
-        {
-            this.functions = ImmutableListMultimap.<QualifiedObjectName, SqlFunction>builder()
-                    .putAll(map.functions)
-                    .putAll(Multimaps.index(functions, function -> function.getSignature().getName()))
-                    .build();
-
-            // Make sure all functions with the same name are aggregations or none of them are
-            for (Map.Entry<QualifiedObjectName, Collection<SqlFunction>> entry : this.functions.asMap().entrySet()) {
-                Collection<SqlFunction> values = entry.getValue();
-                long aggregations = values.stream()
-                        .map(function -> function.getSignature().getKind())
-                        .filter(kind -> kind == AGGREGATE)
-                        .count();
-                checkState(aggregations == 0 || aggregations == values.size(), "'%s' is both an aggregation and a scalar function", entry.getKey());
-            }
-        }
-
-        public List<SqlFunction> list()
-        {
-            return ImmutableList.copyOf(functions.values());
-        }
-
-        public Collection<SqlFunction> get(QualifiedObjectName name)
-        {
-            return functions.get(name);
-        }
-    }
-
     /**
      * TypeSignature but has overridden equals(). Here, we compare exact signature of any underlying distinct
      * types. Some distinct types may have extra information on their lazily loaded parents, and same parent