org.apache.pig.newplan.logical.rules.LimitOptimizer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of pig Show documentation
There is a newer version: 0.17.0
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.pig.newplan.logical.rules;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOCross;
import org.apache.pig.newplan.logical.relational.LODistinct;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOJoin;
import org.apache.pig.newplan.logical.relational.LOLimit;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LOSplit;
import org.apache.pig.newplan.logical.relational.LOSplitOutput;
import org.apache.pig.newplan.logical.relational.LOUnion;
import org.apache.pig.newplan.logical.relational.LogicalPlan;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.optimizer.Rule;
import org.apache.pig.newplan.optimizer.Transformer;

public class LimitOptimizer extends Rule {

    public LimitOptimizer(String name) {
        super(name, false);
    }

    @Override
    protected OperatorPlan buildPattern() {
        LogicalPlan plan = new LogicalPlan();
        LogicalRelationalOperator limit = new LOLimit(plan, 0);
        plan.add(limit);
        return plan;
    }

    @Override
    public Transformer getNewTransformer() {
        return new OptimizeLimitTransformer();
    }

    public class OptimizeLimitTransformer extends Transformer {

        @Override
        public boolean check(OperatorPlan matched) {
            LOLimit limit = (LOLimit) matched.getSources().get(0);
            // Match each foreach.
            List preds = currentPlan.getPredecessors(limit);
            if (preds == null || preds.size() == 0)
                return false;

            Operator pred = preds.get(0);

            // Limit cannot be pushed up
            if (pred instanceof LOCogroup || pred instanceof LOFilter
                    || pred instanceof LOLoad || pred instanceof LOSplit
                    || pred instanceof LODistinct || pred instanceof LOJoin) {
                return false;
            }

            // Limit cannot be pushed in front of ForEach if it has a flatten
            if (pred instanceof LOForEach) {
                LOForEach foreach = (LOForEach) pred;
                LogicalPlan innerPlan = foreach.getInnerPlan();
                Iterator it = innerPlan.getOperators();
                while (it.hasNext()) {
                    Operator op = it.next();
                    if (op instanceof LOGenerate) {
                        LOGenerate gen = (LOGenerate) op;
                        boolean[] flattenFlags = gen.getFlattenFlags();
                        if (flattenFlags != null) {
                            for (boolean flatten : flattenFlags) {
                                if (flatten)
                                    return false;
                            }
                        }
                    }
                }
            }
            return true;
        }

        @Override
        public OperatorPlan reportChanges() {
            return currentPlan;
        }

        @Override
        public void transform(OperatorPlan matched) throws FrontendException {

            LOLimit limit = (LOLimit) matched.getSources().get(0);

            // Find the next foreach operator.
            List preds = currentPlan.getPredecessors(limit);
            Operator pred = preds.get(0);

            if (pred instanceof LOForEach) {
                // We can safely move LOLimit up
                // Get operator before LOForEach
                Operator prepredecessor = currentPlan.getPredecessors(pred)
                    .get(0);
                currentPlan.removeAndReconnect(limit);
                currentPlan.insertBetween(prepredecessor, limit, pred);
            } else if (pred instanceof LOCross || pred instanceof LOUnion) {
                // Limit can be duplicated, and the new instance pushed in front
                // of an operator for the following operators
                // (that is, if you have X->limit, you can transform that to
                // limit->X->limit):
                LOLimit newLimit = null;
                List nodesToProcess = new ArrayList();
                for (Operator prepredecessor : currentPlan
                        .getPredecessors(pred))
                    nodesToProcess.add(prepredecessor);
                for (Operator prepredecessor : nodesToProcess) {
                    if (prepredecessor instanceof LOLimit) {
                        LOLimit l = (LOLimit) prepredecessor;
                        l.setLimit(l.getLimit() < limit.getLimit() ? l
                                .getLimit() : limit.getLimit());
                    } else {
                        newLimit = new LOLimit((LogicalPlan) currentPlan, limit
                                .getLimit());
                        currentPlan.insertBetween(prepredecessor, newLimit, pred);
                    }
                }
            } else if (pred instanceof LOSort) {
                LOSort sort = (LOSort) pred;
                if (sort.getLimit() == -1)
                    sort.setLimit(limit.getLimit());
                else
                    sort.setLimit(sort.getLimit() < limit.getLimit() ? sort
                            .getLimit() : limit.getLimit());

                // remove the limit
                currentPlan.removeAndReconnect(limit);
            } else if (pred instanceof LOLimit) {
                // Limit is merged into another LOLimit
                LOLimit beforeLimit = (LOLimit) pred;
                beforeLimit
                        .setLimit(beforeLimit.getLimit() < limit.getLimit() ? beforeLimit
                                .getLimit()
                                : limit.getLimit());
                // remove the limit
                currentPlan.removeAndReconnect(limit);
            } else if (pred instanceof LOSplitOutput) {
                // Limit and OrderBy (LOSort) can be separated by split
                List grandparants = currentPlan.getPredecessors(pred);
                // After insertion of splitters, any node in the plan can
                // have at most one predecessor
                if (grandparants != null && grandparants.size() != 0
                        && grandparants.get(0) instanceof LOSplit) {
                    List greatGrandparants = currentPlan
                            .getPredecessors(grandparants.get(0));
                    if (greatGrandparants != null
                            && greatGrandparants.size() != 0
                            && greatGrandparants.get(0) instanceof LOSort) {
                        LOSort sort = (LOSort) greatGrandparants.get(0);
                        LOSort newSort = new LOSort(sort.getPlan(), sort
                                .getSortColPlans(), sort.getAscendingCols(),
                                sort.getUserFunc());
                        newSort.setLimit(limit.getLimit());

                        currentPlan.replace(limit, newSort);
                    }
                }
            }
        }
    }
}