io.trino.metadata.InternalFunctionDependencies Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trino-main Show documentation
Trino - Engine
There is a newer version: 465
/*
 * 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.metadata;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.trino.spi.function.CatalogSchemaFunctionName;
import io.trino.spi.function.FunctionDependencies;
import io.trino.spi.function.FunctionNullability;
import io.trino.spi.function.InvocationConvention;
import io.trino.spi.function.OperatorType;
import io.trino.spi.function.ScalarFunctionImplementation;
import io.trino.spi.type.Type;
import io.trino.spi.type.TypeSignature;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static io.trino.metadata.GlobalFunctionCatalog.isBuiltinFunctionName;
import static io.trino.metadata.OperatorNameUtil.isOperatorName;
import static io.trino.metadata.OperatorNameUtil.unmangleOperator;
import static io.trino.spi.function.OperatorType.CAST;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
import static java.util.function.Function.identity;

public class InternalFunctionDependencies
        implements FunctionDependencies
{
    private final BiFunction specialization;
    private final Map types;
    private final Map functions;
    private final Map operators;
    private final Map casts;

    public InternalFunctionDependencies(
            BiFunction specialization,
            Map typeDependencies,
            Collection functionDependencies)
    {
        requireNonNull(specialization, "specialization is null");
        requireNonNull(typeDependencies, "typeDependencies is null");
        requireNonNull(functionDependencies, "functionDependencies is null");

        this.specialization = specialization;
        this.types = ImmutableMap.copyOf(typeDependencies);
        this.functions = functionDependencies.stream()
                .filter(function -> !isOperatorName(function.signature().getName().getFunctionName()))
                .collect(toImmutableMap(FunctionKey::new, identity()));
        this.operators = functionDependencies.stream()
                .filter(InternalFunctionDependencies::isOperator)
                .collect(toImmutableMap(OperatorKey::new, identity()));
        this.casts = functionDependencies.stream()
                .filter(InternalFunctionDependencies::isCast)
                .collect(toImmutableMap(CastKey::new, identity()));
    }

    @Override
    public Type getType(TypeSignature typeSignature)
    {
        // CHAR type does not properly roundtrip, so load directly from metadata and then verify type was declared correctly
        Type type = types.get(typeSignature);
        if (type == null) {
            throw new UndeclaredDependencyException(typeSignature.toString());
        }
        return type;
    }

    @Override
    public FunctionNullability getFunctionNullability(CatalogSchemaFunctionName name, List parameterTypes)
    {
        FunctionKey functionKey = new FunctionKey(name, toTypeSignatures(parameterTypes));
        ResolvedFunction resolvedFunction = functions.get(functionKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(functionKey.toString());
        }
        return resolvedFunction.functionNullability();
    }

    @Override
    public FunctionNullability getOperatorNullability(OperatorType operatorType, List parameterTypes)
    {
        OperatorKey operatorKey = new OperatorKey(operatorType, toTypeSignatures(parameterTypes));
        ResolvedFunction resolvedFunction = operators.get(operatorKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(operatorKey.toString());
        }
        return resolvedFunction.functionNullability();
    }

    @Override
    public FunctionNullability getCastNullability(Type fromType, Type toType)
    {
        CastKey castKey = new CastKey(fromType.getTypeSignature(), toType.getTypeSignature());
        ResolvedFunction resolvedFunction = casts.get(castKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(castKey.toString());
        }
        return resolvedFunction.functionNullability();
    }

    @Override
    public ScalarFunctionImplementation getScalarFunctionImplementation(CatalogSchemaFunctionName name, List parameterTypes, InvocationConvention invocationConvention)
    {
        FunctionKey functionKey = new FunctionKey(name, toTypeSignatures(parameterTypes));
        ResolvedFunction resolvedFunction = functions.get(functionKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(functionKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    @Override
    public ScalarFunctionImplementation getScalarFunctionImplementationSignature(CatalogSchemaFunctionName name, List parameterTypes, InvocationConvention invocationConvention)
    {
        FunctionKey functionKey = new FunctionKey(name, parameterTypes);
        ResolvedFunction resolvedFunction = functions.get(functionKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(functionKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    @Override
    public ScalarFunctionImplementation getOperatorImplementation(OperatorType operatorType, List parameterTypes, InvocationConvention invocationConvention)
    {
        OperatorKey operatorKey = new OperatorKey(operatorType, toTypeSignatures(parameterTypes));
        ResolvedFunction resolvedFunction = operators.get(operatorKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(operatorKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    @Override
    public ScalarFunctionImplementation getOperatorImplementationSignature(OperatorType operatorType, List parameterTypes, InvocationConvention invocationConvention)
    {
        OperatorKey operatorKey = new OperatorKey(operatorType, parameterTypes);
        ResolvedFunction resolvedFunction = operators.get(operatorKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(operatorKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    @Override
    public ScalarFunctionImplementation getCastImplementation(Type fromType, Type toType, InvocationConvention invocationConvention)
    {
        CastKey castKey = new CastKey(fromType.getTypeSignature(), toType.getTypeSignature());
        ResolvedFunction resolvedFunction = casts.get(castKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(castKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    @Override
    public ScalarFunctionImplementation getCastImplementationSignature(TypeSignature fromType, TypeSignature toType, InvocationConvention invocationConvention)
    {
        CastKey castKey = new CastKey(fromType, toType);
        ResolvedFunction resolvedFunction = casts.get(castKey);
        if (resolvedFunction == null) {
            throw new UndeclaredDependencyException(castKey.toString());
        }
        return specialization.apply(resolvedFunction, invocationConvention);
    }

    private static List toTypeSignatures(List types)
    {
        return types.stream()
                .map(Type::getTypeSignature)
                .collect(toImmutableList());
    }

    private static boolean isOperator(ResolvedFunction function)
    {
        CatalogSchemaFunctionName name = function.signature().getName();
        return isBuiltinFunctionName(name) && isOperatorName(name.getFunctionName()) && unmangleOperator(name.getFunctionName()) != CAST;
    }

    private static boolean isCast(ResolvedFunction function)
    {
        CatalogSchemaFunctionName name = function.signature().getName();
        return isBuiltinFunctionName(name) && isOperatorName(name.getFunctionName()) && unmangleOperator(name.getFunctionName()) == CAST;
    }

    public static final class FunctionKey
    {
        private final CatalogSchemaFunctionName name;
        private final List argumentTypes;

        private FunctionKey(ResolvedFunction resolvedFunction)
        {
            name = resolvedFunction.signature().getName();
            argumentTypes = resolvedFunction.signature().getArgumentTypes().stream()
                    .map(Type::getTypeSignature)
                    .collect(toImmutableList());
        }

        private FunctionKey(CatalogSchemaFunctionName name, List argumentTypes)
        {
            this.name = requireNonNull(name, "name is null");
            this.argumentTypes = ImmutableList.copyOf(requireNonNull(argumentTypes, "argumentTypes is null"));
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            FunctionKey that = (FunctionKey) o;
            return Objects.equals(name, that.name) &&
                    Objects.equals(argumentTypes, that.argumentTypes);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(name, argumentTypes);
        }

        @Override
        public String toString()
        {
            return name + argumentTypes.stream()
                    .map(TypeSignature::toString)
                    .collect(Collectors.joining(", ", "(", ")"));
        }
    }

    public static final class OperatorKey
    {
        private final OperatorType operatorType;
        private final List argumentTypes;

        private OperatorKey(ResolvedFunction resolvedFunction)
        {
            operatorType = unmangleOperator(resolvedFunction.signature().getName().getFunctionName());
            argumentTypes = toTypeSignatures(resolvedFunction.signature().getArgumentTypes());
        }

        private OperatorKey(OperatorType operatorType, List argumentTypes)
        {
            this.operatorType = requireNonNull(operatorType, "operatorType is null");
            this.argumentTypes = ImmutableList.copyOf(requireNonNull(argumentTypes, "argumentTypes is null"));
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            OperatorKey that = (OperatorKey) o;
            return operatorType == that.operatorType &&
                    Objects.equals(argumentTypes, that.argumentTypes);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(operatorType, argumentTypes);
        }

        @Override
        public String toString()
        {
            return operatorType + argumentTypes.stream()
                    .map(TypeSignature::toString)
                    .collect(Collectors.joining(", ", "(", ")"));
        }
    }

    private static final class CastKey
    {
        private final TypeSignature fromType;
        private final TypeSignature toType;

        private CastKey(ResolvedFunction resolvedFunction)
        {
            fromType = resolvedFunction.signature().getArgumentTypes().get(0).getTypeSignature();
            toType = resolvedFunction.signature().getReturnType().getTypeSignature();
        }

        private CastKey(TypeSignature fromType, TypeSignature toType)
        {
            this.fromType = fromType;
            this.toType = toType;
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            CastKey that = (CastKey) o;
            return Objects.equals(fromType, that.fromType) &&
                    Objects.equals(toType, that.toType);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(fromType, toType);
        }

        @Override
        public String toString()
        {
            return format("cast(%s, %s)", fromType, toType);
        }
    }
}