io.trino.sql.planner.iterative.rule.GatherPartialTopN 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.sql.planner.iterative.rule;

import com.google.common.collect.ImmutableList;
import io.trino.matching.Capture;
import io.trino.matching.Captures;
import io.trino.matching.Pattern;
import io.trino.sql.planner.Partitioning;
import io.trino.sql.planner.PartitioningScheme;
import io.trino.sql.planner.iterative.Rule;
import io.trino.sql.planner.plan.ExchangeNode;
import io.trino.sql.planner.plan.PlanNode;
import io.trino.sql.planner.plan.TopNNode;

import static io.trino.matching.Capture.newCapture;
import static io.trino.sql.planner.SystemPartitioningHandle.FIXED_ARBITRARY_DISTRIBUTION;
import static io.trino.sql.planner.plan.ExchangeNode.Scope.LOCAL;
import static io.trino.sql.planner.plan.ExchangeNode.Scope.REMOTE;
import static io.trino.sql.planner.plan.ExchangeNode.Type.GATHER;
import static io.trino.sql.planner.plan.ExchangeNode.gatheringExchange;
import static io.trino.sql.planner.plan.ExchangeNode.partitionedExchange;
import static io.trino.sql.planner.plan.Patterns.exchange;
import static io.trino.sql.planner.plan.Patterns.source;
import static io.trino.sql.planner.plan.Patterns.topN;
import static io.trino.sql.planner.plan.TopNNode.Step.PARTIAL;

/**
 * Adds local round-robin and gathering exchange on top of partial TopN to limit the task output size.
 * Replaces plans like:
 *  * exchange(remote)
 *   - topn(partial)
 * 
 * with
 *  *  exchange(remote)
 *      - topn(partial)
 *          - exchange(local, gather)
 *              - topn(partial)
 *                  - local_exchange(round_robin)
 *                      - topn(partial)
 * 
 */
public class GatherPartialTopN
        implements Rule
{
    private static final Capture TOPN = newCapture();
    // the pattern filters for parent and source exchanges are added to avoid infinite recursion in iterative optimizer
    private static final Pattern PATTERN = exchange()
            .matching(GatherPartialTopN::isGatherRemoteExchange)
            .with(source().matching(
                    topN().matching(topN -> topN.getStep().equals(PARTIAL))
                            .with(source().matching(source -> !isGatherLocalExchange(source)))
                            .capturedAs(TOPN)));

    private static boolean isGatherLocalExchange(PlanNode source)
    {
        return source instanceof ExchangeNode exchange
               && exchange.getScope().equals(LOCAL)
               && exchange.getType().equals(GATHER);
    }

    private static boolean isGatherRemoteExchange(ExchangeNode exchangeNode)
    {
        return exchangeNode.getScope().equals(REMOTE)
               && exchangeNode.getType().equals(GATHER)
               // non-empty orderingScheme means it's a merging exchange
               && exchangeNode.getOrderingScheme().isEmpty();
    }

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

    @Override
    public Result apply(ExchangeNode node, Captures captures, Context context)
    {
        TopNNode originalPartialTopN = captures.get(TOPN);

        TopNNode roundRobinTopN = new TopNNode(
                context.getIdAllocator().getNextId(),
                partitionedExchange(
                        context.getIdAllocator().getNextId(),
                        LOCAL,
                        originalPartialTopN,
                        new PartitioningScheme(Partitioning.create(FIXED_ARBITRARY_DISTRIBUTION, ImmutableList.of()), originalPartialTopN.getOutputSymbols())),
                originalPartialTopN.getCount(),
                originalPartialTopN.getOrderingScheme(),
                PARTIAL);

        return Result.ofPlanNode(node.replaceChildren(
                ImmutableList.of(new TopNNode(
                        context.getIdAllocator().getNextId(),
                        gatheringExchange(context.getIdAllocator().getNextId(), LOCAL, roundRobinTopN),
                        originalPartialTopN.getCount(),
                        originalPartialTopN.getOrderingScheme(),
                        PARTIAL))));
    }
}