io.trino.sql.planner.iterative.rule.UnwrapSingleColumnRowInApply 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: 468
/*
 * 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.iterative.rule;

import com.google.common.collect.ImmutableMap;
import io.trino.matching.Captures;
import io.trino.matching.Pattern;
import io.trino.spi.type.RowType;
import io.trino.spi.type.Type;
import io.trino.sql.ir.Expression;
import io.trino.sql.ir.FieldReference;
import io.trino.sql.planner.Symbol;
import io.trino.sql.planner.iterative.Rule;
import io.trino.sql.planner.plan.ApplyNode;
import io.trino.sql.planner.plan.Assignments;
import io.trino.sql.planner.plan.Assignments.Assignment;
import io.trino.sql.planner.plan.ProjectNode;

import java.util.Map;
import java.util.Optional;
import java.util.function.BiFunction;

import static io.trino.sql.planner.plan.Patterns.applyNode;
import static java.util.Objects.requireNonNull;

/**
 * Given x::row(t) and y::row(t), converts assignments of the form
 *
 * {@code x IN (y...) => x[1] IN (y[1]...)}
 *
 * 
and

 *
 * {@code x   (y...) => x[1]   (y[1]...)}
 *
 * 
In particular, it transforms a plan with the following shape:
 *
 * {@code
 * - Apply x IN y
 *   - S [x :: row(T)]
 *   - Q [y :: row(T)]
 * }
 * 
 * into
 *
 * 
{@code
 * - Project (to preserve the outputs of Apply)
 *   - Apply x' IN y'
 *     - Project [x' :: T]
 *         x' = x[1]
 *       - S [x :: row(T)]
 *     - Project [y' :: T]
 *         y' = y[1]
 *       - Q [y :: row(T)]
 * }
 */
public class UnwrapSingleColumnRowInApply
        implements Rule
{
    private static final Pattern PATTERN = applyNode();

    @Override
    public Pattern getPattern()
    {
        return PATTERN;
    }

    @Override
    public Result apply(ApplyNode node, Captures captures, Context context)
    {
        Assignments.Builder inputAssignments = Assignments.builder()
                .putIdentities(node.getInput().getOutputSymbols());

        Assignments.Builder nestedPlanAssignments = Assignments.builder()
                .putIdentities(node.getSubquery().getOutputSymbols());

        boolean applied = false;
        ImmutableMap.Builder applyAssignments = ImmutableMap.builder();
        for (Map.Entry assignment : node.getSubqueryAssignments().entrySet()) {
            Symbol output = assignment.getKey();
            ApplyNode.SetExpression expression = assignment.getValue();

            Optional unwrapped = Optional.empty();
            if (expression instanceof ApplyNode.In predicate) {
                unwrapped = unwrapSingleColumnRow(
                        context,
                        predicate.value().toSymbolReference(),
                        predicate.reference().toSymbolReference(),
                        ApplyNode.In::new);
            }
            else if (expression instanceof ApplyNode.QuantifiedComparison comparison) {
                unwrapped = unwrapSingleColumnRow(
                        context,
                        comparison.value().toSymbolReference(),
                        comparison.reference().toSymbolReference(),
                        (value, list) -> new ApplyNode.QuantifiedComparison(comparison.operator(), comparison.quantifier(), value, list));
            }

            if (unwrapped.isPresent()) {
                applied = true;

                Unwrapping unwrapping = unwrapped.get();
                inputAssignments.add(unwrapping.getInputAssignment());
                nestedPlanAssignments.add(unwrapping.getNestedPlanAssignment());
                applyAssignments.put(output, unwrapping.getExpression());
            }
            else {
                applyAssignments.put(assignment);
            }
        }

        if (!applied) {
            return Result.empty();
        }

        return Result.ofPlanNode(
                new ProjectNode(
                        context.getIdAllocator().getNextId(),
                        new ApplyNode(
                                node.getId(),
                                new ProjectNode(context.getIdAllocator().getNextId(), node.getInput(), inputAssignments.build()),
                                new ProjectNode(context.getIdAllocator().getNextId(), node.getSubquery(), nestedPlanAssignments.build()),
                                applyAssignments.buildOrThrow(),
                                node.getCorrelation(),
                                node.getOriginSubquery()),
                        Assignments.identity(node.getOutputSymbols())));
    }

    private Optional unwrapSingleColumnRow(Context context, Expression value, Expression list, BiFunction function)
    {
        Type type = value.type();
        if (type instanceof RowType rowType) {
            if (rowType.getFields().size() == 1) {
                Type elementType = rowType.getTypeParameters().get(0);

                Symbol valueSymbol = context.getSymbolAllocator().newSymbol("input", elementType);
                Symbol listSymbol = context.getSymbolAllocator().newSymbol("subquery", elementType);

                Assignment inputAssignment = new Assignment(valueSymbol, new FieldReference(value, 0));
                Assignment nestedPlanAssignment = new Assignment(listSymbol, new FieldReference(list, 0));
                ApplyNode.SetExpression comparison = function.apply(valueSymbol, listSymbol);

                return Optional.of(new Unwrapping(comparison, inputAssignment, nestedPlanAssignment));
            }
        }

        return Optional.empty();
    }

    private static class Unwrapping
    {
        private final ApplyNode.SetExpression expression;
        private final Assignment inputAssignment;
        private final Assignment nestedPlanAssignment;

        public Unwrapping(ApplyNode.SetExpression expression, Assignment inputAssignment, Assignment nestedPlanAssignment)
        {
            this.expression = requireNonNull(expression, "expression is null");
            this.inputAssignment = requireNonNull(inputAssignment, "inputAssignment is null");
            this.nestedPlanAssignment = requireNonNull(nestedPlanAssignment, "nestedPlanAssignment is null");
        }

        public ApplyNode.SetExpression getExpression()
        {
            return expression;
        }

        public Assignment getInputAssignment()
        {
            return inputAssignment;
        }

        public Assignment getNestedPlanAssignment()
        {
            return nestedPlanAssignment;
        }
    }
}