All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.trino.sql.planner.Partitioning Maven / Gradle / Ivy

/*
 * 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.sql.planner;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.errorprone.annotations.DoNotCall;
import com.google.errorprone.annotations.Immutable;
import io.trino.Session;
import io.trino.metadata.Metadata;
import io.trino.spi.predicate.NullableValue;
import io.trino.sql.ir.Expression;
import io.trino.sql.ir.Reference;

import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Verify.verify;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static java.util.Objects.requireNonNull;

@Immutable
public final class Partitioning
{
    private final PartitioningHandle handle;
    private final List arguments;

    private Partitioning(PartitioningHandle handle, List arguments)
    {
        this.handle = requireNonNull(handle, "handle is null");
        this.arguments = ImmutableList.copyOf(requireNonNull(arguments, "arguments is null"));
    }

    public static Partitioning create(PartitioningHandle handle, List columns)
    {
        return new Partitioning(handle, columns.stream()
                .map(Symbol::toSymbolReference)
                .map(ArgumentBinding::expressionBinding)
                .collect(toImmutableList()));
    }

    @JsonCreator
    @DoNotCall // For JSON deserialization only
    public static Partitioning jsonCreate(
            @JsonProperty("handle") PartitioningHandle handle,
            @JsonProperty("arguments") List arguments)
    {
        return new Partitioning(handle, arguments);
    }

    @JsonProperty
    public PartitioningHandle getHandle()
    {
        return handle;
    }

    @JsonProperty
    public List getArguments()
    {
        return arguments;
    }

    public Set getColumns()
    {
        return arguments.stream()
                .filter(ArgumentBinding::isVariable)
                .map(ArgumentBinding::getColumn)
                .collect(toImmutableSet());
    }

    public boolean isCompatibleWith(
            Partitioning right,
            Metadata metadata,
            Session session)
    {
        if (!handle.equals(right.handle) && metadata.getCommonPartitioning(session, handle, right.handle).isEmpty()) {
            return false;
        }

        return arguments.equals(right.arguments);
    }

    public boolean isCompatibleWith(
            Partitioning right,
            Function> leftToRightMappings,
            Function> leftConstantMapping,
            Function> rightConstantMapping,
            Metadata metadata,
            Session session)
    {
        if (!handle.equals(right.handle) && metadata.getCommonPartitioning(session, handle, right.handle).isEmpty()) {
            return false;
        }

        if (arguments.size() != right.arguments.size()) {
            return false;
        }

        for (int i = 0; i < arguments.size(); i++) {
            ArgumentBinding leftArgument = arguments.get(i);
            ArgumentBinding rightArgument = right.arguments.get(i);

            if (!isPartitionedWith(leftArgument, leftConstantMapping, rightArgument, rightConstantMapping, leftToRightMappings)) {
                return false;
            }
        }
        return true;
    }

    private static boolean isPartitionedWith(
            ArgumentBinding leftArgument,
            Function> leftConstantMapping,
            ArgumentBinding rightArgument,
            Function> rightConstantMapping,
            Function> leftToRightMappings)
    {
        if (leftArgument.isVariable()) {
            if (rightArgument.isVariable()) {
                // variable == variable
                Set mappedColumns = leftToRightMappings.apply(leftArgument.getColumn());
                return mappedColumns.contains(rightArgument.getColumn());
            }
            // variable == constant
            // Normally, this would be a false condition, but if we happen to have an external
            // mapping from the symbol to a constant value and that constant value matches the
            // right value, then we are co-partitioned.
            Optional leftConstant = leftConstantMapping.apply(leftArgument.getColumn());
            return leftConstant.isPresent() && leftConstant.get().equals(rightArgument.getConstant());
        }
        if (rightArgument.isConstant()) {
            // constant == constant
            return leftArgument.getConstant().equals(rightArgument.getConstant());
        }
        // constant == variable
        Optional rightConstant = rightConstantMapping.apply(rightArgument.getColumn());
        return rightConstant.isPresent() && rightConstant.get().equals(leftArgument.getConstant());
    }

    public boolean isPartitionedOn(Collection columns, Set knownConstants)
    {
        for (ArgumentBinding argument : arguments) {
            // partitioned on (k_1, k_2, ..., k_n) => partitioned on (k_1, k_2, ..., k_n, k_n+1, ...)
            // can safely ignore all constant columns when comparing partition properties
            if (argument.isConstant()) {
                continue;
            }
            if (!argument.isVariable()) {
                return false;
            }
            if (!knownConstants.contains(argument.getColumn()) && !columns.contains(argument.getColumn())) {
                return false;
            }
        }
        return true;
    }

    public boolean isPartitionedOnExactly(Collection columns, Set knownConstants)
    {
        Set toCheck = new HashSet<>();
        for (ArgumentBinding argument : arguments) {
            // partitioned on (k_1, k_2, ..., k_n) => partitioned on (k_1, k_2, ..., k_n, k_n+1, ...)
            // can safely ignore all constant columns when comparing partition properties
            if (argument.isConstant()) {
                continue;
            }
            if (!argument.isVariable()) {
                return false;
            }
            if (knownConstants.contains(argument.getColumn())) {
                continue;
            }
            toCheck.add(argument.getColumn());
        }
        return ImmutableSet.copyOf(columns).equals(toCheck);
    }

    public boolean isEffectivelySinglePartition(Set knownConstants)
    {
        return isPartitionedOn(ImmutableSet.of(), knownConstants);
    }

    public Partitioning translate(Function translator)
    {
        return new Partitioning(handle, arguments.stream()
                .map(argument -> argument.translate(translator))
                .collect(toImmutableList()));
    }

    public Optional translate(Translator translator)
    {
        ImmutableList.Builder newArguments = ImmutableList.builder();
        for (ArgumentBinding argument : arguments) {
            Optional newArgument = argument.translate(translator);
            if (newArgument.isEmpty()) {
                return Optional.empty();
            }
            newArguments.add(newArgument.get());
        }

        return Optional.of(new Partitioning(handle, newArguments.build()));
    }

    public Partitioning withAlternativePartitioningHandle(PartitioningHandle partitioningHandle)
    {
        return new Partitioning(partitioningHandle, this.arguments);
    }

    @Override
    public int hashCode()
    {
        return Objects.hash(handle, arguments);
    }

    @Override
    public boolean equals(Object obj)
    {
        if (this == obj) {
            return true;
        }
        if (obj == null || getClass() != obj.getClass()) {
            return false;
        }
        Partitioning other = (Partitioning) obj;
        return Objects.equals(this.handle, other.handle) &&
                Objects.equals(this.arguments, other.arguments);
    }

    @Override
    public String toString()
    {
        return toStringHelper(this)
                .add("handle", handle)
                .add("arguments", arguments)
                .toString();
    }

    @Immutable
    public static final class Translator
    {
        private final Function> columnTranslator;
        private final Function> constantTranslator;
        private final Function> expressionTranslator;

        public Translator(
                Function> columnTranslator,
                Function> constantTranslator,
                Function> expressionTranslator)
        {
            this.columnTranslator = requireNonNull(columnTranslator, "columnTranslator is null");
            this.constantTranslator = requireNonNull(constantTranslator, "constantTranslator is null");
            this.expressionTranslator = requireNonNull(expressionTranslator, "expressionTranslator is null");
        }
    }

    @Immutable
    public static final class ArgumentBinding
    {
        private final Expression expression;
        private final NullableValue constant;

        @JsonCreator
        public ArgumentBinding(
                @JsonProperty("expression") Expression expression,
                @JsonProperty("constant") NullableValue constant)
        {
            this.expression = expression;
            this.constant = constant;
            checkArgument((expression == null) != (constant == null), "Either expression or constant must be set");
        }

        public static ArgumentBinding expressionBinding(Expression expression)
        {
            return new ArgumentBinding(requireNonNull(expression, "expression is null"), null);
        }

        public static ArgumentBinding constantBinding(NullableValue constant)
        {
            return new ArgumentBinding(null, requireNonNull(constant, "constant is null"));
        }

        public boolean isConstant()
        {
            return constant != null;
        }

        public boolean isVariable()
        {
            return expression instanceof Reference;
        }

        public Symbol getColumn()
        {
            verify(expression instanceof Reference, "Expect the expression to be a SymbolReference");
            return Symbol.from(expression);
        }

        @JsonProperty
        public Expression getExpression()
        {
            return expression;
        }

        @JsonProperty
        public NullableValue getConstant()
        {
            return constant;
        }

        public ArgumentBinding translate(Function translator)
        {
            if (isConstant()) {
                return this;
            }
            return expressionBinding(translator.apply(Symbol.from(expression)).toSymbolReference());
        }

        public Optional translate(Translator translator)
        {
            if (isConstant()) {
                return Optional.of(this);
            }

            if (!isVariable()) {
                return translator.expressionTranslator.apply(expression)
                        .map(Symbol::toSymbolReference)
                        .map(ArgumentBinding::expressionBinding);
            }

            Optional newColumn = translator.columnTranslator.apply(Symbol.from(expression))
                    .map(Symbol::toSymbolReference)
                    .map(ArgumentBinding::expressionBinding);
            if (newColumn.isPresent()) {
                return newColumn;
            }
            // As a last resort, check for a constant mapping for the symbol
            // Note: this MUST be last because we want to favor the symbol representation
            // as it makes further optimizations possible.
            return translator.constantTranslator.apply(Symbol.from(expression))
                    .map(ArgumentBinding::constantBinding);
        }

        @Override
        public String toString()
        {
            if (constant != null) {
                return constant.toString();
            }

            return expression.toString();
        }

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

        @Override
        public int hashCode()
        {
            return Objects.hash(expression, constant);
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy