/*
* 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.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
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.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<Operator> 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 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<Operator> 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<Operator> 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);
List<Operator> softPrepredecessors=null;
// get a clone of softPrepredecessors to avoid ConcurrentModificationException
if (currentPlan.getSoftLinkPredecessors(limit)!=null) {
softPrepredecessors=new ArrayList<Operator>(
currentPlan.getSoftLinkPredecessors(limit));
}
if (softPrepredecessors!=null) {
for (Operator op : softPrepredecessors) {
currentPlan.removeSoftLink(op, limit);
}
}
currentPlan.removeAndReconnect(limit);
currentPlan.insertBetween(prepredecessor, limit, pred);
if (softPrepredecessors!=null) {
for (Operator op : softPrepredecessors) {
currentPlan.createSoftLink(op, limit);
}
}
} else if (limit.getLimitPlan() == null) {
// TODO selectively enable optimizations for variable limit
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<Operator> nodesToProcess = new ArrayList<Operator>();
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 LOLoad) {
// Push limit to load
LOLoad load = (LOLoad) pred;
if (load.getLimit() == -1)
load.setLimit(limit.getLimit());
else
load.setLimit(load.getLimit() < limit.getLimit() ? load
.getLimit() : limit.getLimit());
} 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<Operator> 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<Operator> 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 = LOSort.createCopy(sort);
newSort.setLimit(limit.getLimit());
currentPlan.replace(limit, newSort);
}
}
}
}
}
}
}