Support multiple annotation based definitions of scalar function

For datatypes with multiple binary representations we need to define
same function multiple times depending on actal types of parameters.

E.g.
  add (Slice a, Slice b)
and
  add (long a, long b)

represent + operator for DECIMAL type. First version for longer DECIMAL
values (represented internally as Slice) and second, faster version for
shorter decimals (represented internally as long).

This patch allows multiple versions of same scalar udf/operator to be
defined using annotation syntax. Functions with same name are grouped
together and exposed as single ParametricFunction which internally does
dispatching to proper MethodHandle based on actual parameter types at
runtime.

Author: losipiuk

presto-main/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java

@@ -14,7 +14,6 @@
 package com.facebook.presto.metadata;
 
 import com.facebook.presto.operator.Description;
-import com.facebook.presto.operator.aggregation.GenericAggregationFunctionFactory;
 import com.facebook.presto.operator.aggregation.InternalAggregationFunction;
 import com.facebook.presto.operator.scalar.JsonPath;
 import com.facebook.presto.operator.scalar.ScalarFunction;
@@ -46,19 +45,26 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static com.facebook.presto.metadata.FunctionKind.SCALAR;
 import static com.facebook.presto.metadata.FunctionKind.WINDOW;
 import static com.facebook.presto.metadata.Signature.typeParameter;
+import static com.facebook.presto.operator.aggregation.GenericAggregationFunctionFactory.fromAggregationDefinition;
 import static com.facebook.presto.spi.type.BigintType.BIGINT;
 import static com.facebook.presto.spi.type.TypeSignature.parseTypeSignature;
 import static com.google.common.base.CaseFormat.LOWER_CAMEL;
 import static com.google.common.base.CaseFormat.LOWER_UNDERSCORE;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static java.lang.invoke.MethodHandles.lookup;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.toList;
 
 public class FunctionListBuilder
 {
@@ -69,7 +75,35 @@ public class FunctionListBuilder
             Regex.class,
             JsonPath.class);
 
-    private final List<SqlFunction> functions = new ArrayList<>();
+    private static final class FunctionWithMethodHandle
+    {
+        private final SqlFunction function;
+        private final Optional<MethodHandle> methodHandle;
+
+        public FunctionWithMethodHandle(SqlFunction function)
+        {
+            this.function = function;
+            this.methodHandle = Optional.empty();
+        }
+
+        public FunctionWithMethodHandle(SqlFunction function, MethodHandle methodHandle)
+        {
+            this.function = function;
+            this.methodHandle = Optional.of(methodHandle);
+        }
+
+        public Optional<MethodHandle> getMethodHandle()
+        {
+            return methodHandle;
+        }
+
+        public SqlFunction getFunction()
+        {
+            return function;
+        }
+    }
+
+    private final List<FunctionWithMethodHandle> functions = new ArrayList<>();
     private final TypeManager typeManager;
 
     public FunctionListBuilder(TypeManager typeManager)
@@ -83,14 +117,14 @@ public FunctionListBuilder window(String name, Type returnType, List<? extends T
                 new Signature(name, WINDOW, returnType.getTypeSignature(), Lists.transform(ImmutableList.copyOf(argumentTypes), Type::getTypeSignature)),
                 functionClass);
 
-        functions.add(new SqlWindowFunction(windowFunctionSupplier));
+        addFunction(new SqlWindowFunction(windowFunctionSupplier));
         return this;
     }
 
     public FunctionListBuilder window(String name, Class<? extends ValueWindowFunction> clazz, String typeVariable, String... argumentTypes)
     {
         Signature signature = new Signature(name, WINDOW, ImmutableList.of(typeParameter(typeVariable)), typeVariable, ImmutableList.copyOf(argumentTypes), false);
-        functions.add(new SqlWindowFunction(new ReflectionWindowFunctionSupplier<>(signature, clazz)));
+        addFunction(new SqlWindowFunction(new ReflectionWindowFunctionSupplier<>(signature, clazz)));
         return this;
     }
 
@@ -108,29 +142,47 @@ public FunctionListBuilder aggregate(InternalAggregationFunction function)
         name = name.toLowerCase(ENGLISH);
 
         String description = getDescription(function.getClass());
-        functions.add(SqlAggregationFunction.create(name, description, function));
+        addFunction(SqlAggregationFunction.create(name, description, function));
         return this;
     }
 
     public FunctionListBuilder aggregate(Class<?> aggregationDefinition)
     {
-        functions.addAll(GenericAggregationFunctionFactory.fromAggregationDefinition(aggregationDefinition, typeManager).listFunctions());
+        fromAggregationDefinition(aggregationDefinition, typeManager).listFunctions().forEach(this::addFunction);
         return this;
     }
 
     public FunctionListBuilder scalar(Signature signature, MethodHandle function, boolean deterministic, String description, boolean hidden, boolean nullable, List<Boolean> nullableArguments)
     {
-        functions.add(SqlScalarFunction.create(signature, description, hidden, function, deterministic, nullable, nullableArguments));
+        addFunction(SqlScalarFunction.create(signature, description, hidden, function, deterministic, nullable, nullableArguments), function);
         return this;
     }
 
     private FunctionListBuilder operator(OperatorType operatorType, TypeSignature returnType, List<TypeSignature> parameterTypes, MethodHandle function, boolean nullable, List<Boolean> nullableArguments)
     {
-        functions.add(SqlOperator.create(operatorType, parameterTypes, returnType, function, nullable, nullableArguments));
+        addFunction(SqlOperator.create(operatorType, parameterTypes, returnType, function, nullable, nullableArguments), function);
         return this;
     }
 
     public FunctionListBuilder scalar(Class<?> clazz)
+    {
+        FunctionListBuilder localFunctionListBuilder = new FunctionListBuilder(typeManager);
+        localFunctionListBuilder.processScalarsInClass(clazz);
+        functions.addAll(localFunctionListBuilder.functions);
+        return this;
+    }
+
+    private void addFunction(SqlFunction function)
+    {
+        functions.add(new FunctionWithMethodHandle(function));
+    }
+
+    private void addFunction(SqlFunction function, MethodHandle methodHandle)
+    {
+        functions.add(new FunctionWithMethodHandle(function, methodHandle));
+    }
+
+    private FunctionListBuilder processScalarsInClass(Class<?> clazz)
     {
         try {
             boolean foundOne = false;
@@ -143,9 +195,61 @@ public FunctionListBuilder scalar(Class<?> clazz)
         catch (IllegalAccessException e) {
             throw Throwables.propagate(e);
         }
+        groupMatchingScalars();
         return this;
     }
 
+    private FunctionListBuilder groupMatchingScalars()
+    {
+        List<FunctionWithMethodHandle> newFunctions = groupMatchingScalars(functions);
+        functions.clear();
+        functions.addAll(newFunctions);
+        return this;
+    }
+
+    private List<FunctionWithMethodHandle> groupMatchingScalars(List<FunctionWithMethodHandle> inputFunctions)
+    {
+        inputFunctions.forEach(f -> checkArgument(f.getMethodHandle().isPresent(), "expected function with method handle: %s", f));
+        inputFunctions.forEach(f -> checkArgument(f.getFunction() instanceof SqlScalarFunction, "expected scalar function: %s", f));
+        Map<Signature, List<FunctionWithMethodHandle>> groupedFunctions = inputFunctions.stream()
+                .collect(Collectors.groupingBy(f -> f.getFunction().getSignature()));
+        List<FunctionWithMethodHandle> resultFunctions = new ArrayList<>();
+        for (Map.Entry<Signature, List<FunctionWithMethodHandle>> entry : groupedFunctions.entrySet()) {
+            List<FunctionWithMethodHandle> functionsGroup = entry.getValue();
+            if (functionsGroup.size() == 1) {
+                resultFunctions.add(getOnlyElement(functionsGroup));
+            }
+            else {
+                resultFunctions.add(new FunctionWithMethodHandle(buildGroupingScalarWrapper(functionsGroup)));
+            }
+        }
+        return resultFunctions;
+    }
+
+    private SqlScalarFunction buildGroupingScalarWrapper(List<FunctionWithMethodHandle> functionsGroup)
+    {
+        checkArgument(functionsGroup.size() > 1, "functions group must have multiple elements");
+        SqlFunction masterFunction = functionsGroup.get(0).getFunction();
+        functionsGroup.forEach(f -> checkMatchingScalarsConsistent(masterFunction, f.getFunction()));
+
+        SqlScalarFunctionBuilder wrapperBuilder = SqlScalarFunction.builder()
+                .signature(masterFunction.getSignature())
+                .deterministic(masterFunction.isDeterministic())
+                .hidden(masterFunction.isHidden())
+                .description(masterFunction.getDescription());
+
+        functionsGroup.forEach(f -> wrapperBuilder.method((SqlScalarFunction) f.getFunction(), f.getMethodHandle().get()));
+        return wrapperBuilder.build();
+    }
+
+    private void checkMatchingScalarsConsistent(SqlFunction master, SqlFunction canditate)
+    {
+        checkArgument(canditate.getSignature().equals(master.getSignature()), "signature mismatch; %s vs. %s", master, canditate);
+        checkArgument(canditate.isHidden() == master.isHidden(), "hidden flag mismatch; %s vs. %s", master, canditate);
+        checkArgument(canditate.isDeterministic() == master.isDeterministic(), "deterministic flag mismatch; %s vs. %s", master, canditate);
+        checkArgument(Objects.equals(canditate.getDescription(), master.getDescription()), "description mismatch, %s vs. %s", master, canditate);
+    }
+
     public FunctionListBuilder functions(SqlFunction... sqlFunctions)
     {
         for (SqlFunction sqlFunction : sqlFunctions) {
@@ -157,7 +261,7 @@ public FunctionListBuilder functions(SqlFunction... sqlFunctions)
     public FunctionListBuilder function(SqlFunction sqlFunction)
     {
         requireNonNull(sqlFunction, "parametricFunction is null");
-        functions.add(sqlFunction);
+        addFunction(sqlFunction);
         return this;
     }
 
@@ -349,6 +453,6 @@ private static List<Class<?>> getParameterTypes(Class<?>... types)
 
     public List<SqlFunction> getFunctions()
     {
-        return ImmutableList.copyOf(functions);
+        return ImmutableList.copyOf(functions.stream().map(FunctionWithMethodHandle::getFunction).collect(toList()));
     }
 }

presto-main/src/main/java/com/facebook/presto/metadata/PolymorphicScalarFunction.java

@@ -15,15 +15,15 @@
 
 import com.facebook.presto.metadata.SqlScalarFunctionBuilder.MethodsGroup;
 import com.facebook.presto.metadata.SqlScalarFunctionBuilder.SpecializeContext;
+import com.facebook.presto.metadata.SqlScalarFunctionBuilder.TargetMethodDelegate;
 import com.facebook.presto.operator.scalar.ScalarFunctionImplementation;
 import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.spi.type.TypeManager;
 import com.facebook.presto.spi.type.TypeSignature;
-import com.facebook.presto.util.Reflection;
 
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
-import java.lang.reflect.Method;
+import java.lang.invoke.MethodType;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -90,10 +90,10 @@ public ScalarFunctionImplementation specialize(Map<String, Type> types, List<Typ
         Type resolvedReturnType = resolveReturnType(types, typeManager, calculatedReturnType);
         SpecializeContext context = new SpecializeContext(types, filterPresentLiterals(literalParameters), resolvedParameterTypes, resolvedReturnType, typeManager, functionRegistry);
 
-        Optional<Method> matchingMethod = Optional.empty();
+        Optional<TargetMethodDelegate> matchingMethod = Optional.empty();
         Optional<MethodsGroup> matchingMethodsGroup = Optional.empty();
         for (MethodsGroup candidateMethodsGroup : methodsGroups) {
-            for (Method candidateMethod : candidateMethodsGroup.getMethods()) {
+            for (TargetMethodDelegate candidateMethod : candidateMethodsGroup.getMethods()) {
                 if (matchesParameterAndReturnTypes(candidateMethod, resolvedParameterTypes, resolvedReturnType) &&
                         predicateIsTrue(candidateMethodsGroup, context)) {
                     if (matchingMethod.isPresent()) {
@@ -111,10 +111,14 @@ public ScalarFunctionImplementation specialize(Map<String, Type> types, List<Typ
         }
         checkState(matchingMethod.isPresent(), "no matching method for parameter types %s", parameterTypes);
 
-        List<Object> extraParameters = computeExtraParameters(matchingMethodsGroup.get(), context);
-        MethodHandle matchingMethodHandle = applyExtraParameters(matchingMethod.get(), extraParameters);
-
-        return new ScalarFunctionImplementation(nullableResult, nullableArguments, matchingMethodHandle, deterministic);
+        if (matchingMethod.get().getScalarFunctionDelegate().isPresent()) {
+            return matchingMethod.get().getScalarFunctionDelegate().get().specialize(types, parameterTypes, typeManager, functionRegistry);
+        }
+        else {
+            List<Object> extraParameters = computeExtraParameters(matchingMethodsGroup.get(), context);
+            MethodHandle matchingMethodHandle = applyExtraParameters(matchingMethod.get(), extraParameters);
+            return new ScalarFunctionImplementation(nullableResult, nullableArguments, matchingMethodHandle, deterministic);
+        }
     }
 
     private Type resolveReturnType(Map<String, Type> types, TypeManager typeManager, TypeSignature calculatedReturnType)
@@ -129,19 +133,19 @@ private Type resolveReturnType(Map<String, Type> types, TypeManager typeManager,
         return resolvedReturnType;
     }
 
-    private boolean matchesParameterAndReturnTypes(Method method, List<Type> resolvedTypes, Type returnType)
+    private boolean matchesParameterAndReturnTypes(TargetMethodDelegate method, List<Type> resolvedTypes, Type returnType)
     {
-        checkState(method.getParameterCount() >= resolvedTypes.size(),
-                "method %s has not enough arguments: %s (should have at least %s)", method.getName(), method.getParameterCount(), resolvedTypes.size());
+        MethodType methodHandleType = method.getMethodHandle().type();
+        checkState(methodHandleType.parameterCount() >= resolvedTypes.size(),
+                "method %s has not enough arguments: %s (should have at least %s)", method.getName(), methodHandleType.parameterCount(), resolvedTypes.size());
 
-        Class<?>[] methodParameterJavaTypes = method.getParameterTypes();
+        List<Class<?>> methodParameterJavaTypes = methodHandleType.parameterList();
         for (int i = 0; i < resolvedTypes.size(); ++i) {
-            if (!methodParameterJavaTypes[i].equals(resolvedTypes.get(i).getJavaType())) {
+            if (!methodParameterJavaTypes.get(i).equals(resolvedTypes.get(i).getJavaType())) {
                 return false;
             }
         }
-
-        return method.getReturnType().equals(returnType.getJavaType());
+        return methodHandleType.returnType().equals(returnType.getJavaType());
     }
 
     private boolean onlyFirstMatchedMethodHasPredicate(MethodsGroup matchingMethodsGroup, MethodsGroup methodsGroup)
@@ -166,16 +170,15 @@ private Map<String, Long> filterPresentLiterals(Map<String, OptionalLong> boundL
                 .collect(toMap(entry -> entry.getKey().toLowerCase(US), entry -> entry.getValue().getAsLong()));
     }
 
-    private MethodHandle applyExtraParameters(Method matchingMethod, List<Object> extraParameters)
+    private MethodHandle applyExtraParameters(TargetMethodDelegate matchingMethod, List<Object> extraParameters)
     {
         Signature signature = getSignature();
         int expectedNumberOfArguments = signature.getArgumentTypes().size() + extraParameters.size();
-        int matchingMethodParameterCount = matchingMethod.getParameterCount();
+        int matchingMethodParameterCount = matchingMethod.getMethodHandle().type().parameterCount();
         checkState(matchingMethodParameterCount == expectedNumberOfArguments,
                 "method %s has invalid number of arguments: %s (should have %s)", matchingMethod.getName(), matchingMethodParameterCount, expectedNumberOfArguments);
 
-        MethodHandle matchingMethodHandle = Reflection.methodHandle(matchingMethod);
-        matchingMethodHandle = MethodHandles.insertArguments(matchingMethodHandle, signature.getArgumentTypes().size(), extraParameters.toArray());
+        MethodHandle matchingMethodHandle = MethodHandles.insertArguments(matchingMethod.getMethodHandle(), signature.getArgumentTypes().size(), extraParameters.toArray());
         return matchingMethodHandle;
     }
 }

presto-main/src/main/java/com/facebook/presto/metadata/SqlScalarFunction.java

@@ -78,6 +78,10 @@ public static SqlScalarFunctionBuilder builder(Class<?> clazz)
         return new SqlScalarFunctionBuilder(clazz);
     }
 
+    public static SqlScalarFunctionBuilder builder()
+    {
+        return new SqlScalarFunctionBuilder();
+    }
 
     private static class SimpleSqlScalarFunction
             extends SqlScalarFunction