io.trino.operator.scalar.ChoicesSpecializedSqlScalarFunction Maven / Gradle / Ivy

Go to download
/*
 * 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 io.trino.operator.scalar;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import io.trino.spi.TrinoException;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.function.BoundSignature;
import io.trino.spi.function.InvocationConvention;
import io.trino.spi.function.InvocationConvention.InvocationArgumentConvention;
import io.trino.spi.function.InvocationConvention.InvocationReturnConvention;
import io.trino.spi.function.ScalarFunctionAdapter;
import io.trino.spi.function.ScalarFunctionImplementation;

import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkArgument;
import static io.trino.spi.StandardErrorCode.FUNCTION_NOT_FOUND;
import static java.lang.String.format;
import static java.util.Comparator.comparingInt;
import static java.util.Objects.requireNonNull;

public final class ChoicesSpecializedSqlScalarFunction
        implements SpecializedSqlScalarFunction
{
    private final BoundSignature boundSignature;
    private final List choices;

    public ChoicesSpecializedSqlScalarFunction(
            BoundSignature boundSignature,
            InvocationReturnConvention returnConvention,
            List argumentConventions,
            MethodHandle methodHandle)
    {
        this(boundSignature, returnConvention, argumentConventions, ImmutableList.of(), methodHandle, Optional.empty());
    }

    public ChoicesSpecializedSqlScalarFunction(
            BoundSignature boundSignature,
            InvocationReturnConvention returnConvention,
            List argumentConventions,
            MethodHandle methodHandle,
            Optional instanceFactory)
    {
        this(boundSignature, returnConvention, argumentConventions, ImmutableList.of(), methodHandle, instanceFactory);
    }

    public ChoicesSpecializedSqlScalarFunction(
            BoundSignature boundSignature,
            InvocationReturnConvention returnConvention,
            List argumentConventions,
            List> lambdaInterfaces,
            MethodHandle methodHandle,
            Optional instanceFactory)
    {
        this(boundSignature, ImmutableList.of(new ScalarImplementationChoice(returnConvention, argumentConventions, lambdaInterfaces, methodHandle, instanceFactory)));
    }

    /**
     * Creates a ScalarFunctionImplementation consisting of one or more choices.
     * 
     * All choices must have the same SQL signature, and are equivalent in what they do.
     * The first choice is the default choice, which is the one used for legacy access methods.
     * The default choice must be usable under any context. (e.g. it must not use BLOCK_POSITION convention.)
     *
     * @param choices the list of choices, ordered from generic to specific
     */
    public ChoicesSpecializedSqlScalarFunction(BoundSignature boundSignature, List choices)
    {
        this.boundSignature = boundSignature;
        checkArgument(!choices.isEmpty(), "choices is an empty list");
        this.choices = ImmutableList.copyOf(choices);
    }

    @VisibleForTesting
    public List getChoices()
    {
        return choices;
    }

    @Override
    public ScalarFunctionImplementation getScalarFunctionImplementation(InvocationConvention invocationConvention)
    {
        List choices = new ArrayList<>();
        for (ScalarImplementationChoice choice : this.choices) {
            InvocationConvention callingConvention = choice.getInvocationConvention();
            if (ScalarFunctionAdapter.canAdapt(callingConvention, invocationConvention)) {
                choices.add(choice);
            }
        }
        if (choices.isEmpty()) {
            throw new TrinoException(FUNCTION_NOT_FOUND,
                    format("Function implementation for (%s) cannot be adapted to convention (%s)", boundSignature, invocationConvention));
        }

        ScalarImplementationChoice bestChoice = Collections.max(choices, comparingInt(ScalarImplementationChoice::getScore));
        MethodHandle methodHandle = ScalarFunctionAdapter.adapt(
                bestChoice.getMethodHandle(),
                boundSignature.getReturnType(),
                boundSignature.getArgumentTypes(),
                bestChoice.getInvocationConvention(),
                invocationConvention);
        ScalarFunctionImplementation.Builder builder = ScalarFunctionImplementation.builder()
                .methodHandle(methodHandle);
        bestChoice.getInstanceFactory().ifPresent(builder::instanceFactory);
        builder.lambdaInterfaces(bestChoice.getLambdaInterfaces());
        return builder.build();
    }

    public static class ScalarImplementationChoice
    {
        private final MethodHandle methodHandle;
        private final Optional instanceFactory;
        private final InvocationConvention invocationConvention;
        private final List> lambdaInterfaces;
        private final int score;

        public ScalarImplementationChoice(
                InvocationReturnConvention returnConvention,
                List argumentConventions,
                List> lambdaInterfaces,
                MethodHandle methodHandle,
                Optional instanceFactory)
        {
            this.methodHandle = requireNonNull(methodHandle, "methodHandle is null");
            this.instanceFactory = requireNonNull(instanceFactory, "instanceFactory is null");

            if (instanceFactory.isPresent()) {
                Class instanceType = instanceFactory.get().type().returnType();
                checkArgument(instanceFactory.get().type().parameterList().isEmpty(), "instanceFactory should have no parameter");
                checkArgument(instanceType.equals(methodHandle.type().parameterType(0)), "methodHandle is not an instance method");
            }

            List> parameterList = methodHandle.type().parameterList();
            boolean hasSession = false;
            if (parameterList.contains(ConnectorSession.class)) {
                checkArgument(parameterList.stream().filter(ConnectorSession.class::equals).count() == 1, "function implementation should have exactly one ConnectorSession parameter");
                if (instanceFactory.isEmpty()) {
                    checkArgument(parameterList.get(0) == ConnectorSession.class, "ConnectorSession must be the first argument when instanceFactory is not present");
                }
                else {
                    checkArgument(parameterList.get(1) == ConnectorSession.class, "ConnectorSession must be the second argument when instanceFactory is present");
                }
                hasSession = true;
            }

            this.lambdaInterfaces = ImmutableList.copyOf(requireNonNull(lambdaInterfaces, "lambdaInterfaces is null"));
            invocationConvention = new InvocationConvention(
                    argumentConventions,
                    returnConvention,
                    hasSession,
                    instanceFactory.isPresent());
            checkArgument(lambdaInterfaces.size() <= argumentConventions.size());

            score = computeScore(invocationConvention);
        }

        public MethodHandle getMethodHandle()
        {
            return methodHandle;
        }

        public Optional getInstanceFactory()
        {
            return instanceFactory;
        }

        public List> getLambdaInterfaces()
        {
            return lambdaInterfaces;
        }

        public InvocationConvention getInvocationConvention()
        {
            return invocationConvention;
        }

        public int getScore()
        {
            return score;
        }

        private static int computeScore(InvocationConvention callingConvention)
        {
            int score = 0;
            for (InvocationArgumentConvention argument : callingConvention.getArgumentConventions()) {
                switch (argument) {
                    case NULL_FLAG:
                        score += 1;
                        break;
                    case BLOCK_POSITION_NOT_NULL:
                    case BLOCK_POSITION:
                        score += 1000;
                        break;
                    case VALUE_BLOCK_POSITION_NOT_NULL:
                    case VALUE_BLOCK_POSITION:
                        score += 2000;
                        break;
                    case IN_OUT:
                        score += 10_000;
                        break;
                    default:
                        break;
                }
            }
            return score;
        }
    }
}