/**
* 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.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.OperatorSubPlan;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.ProjectExpression;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOInnerLoad;
import org.apache.pig.newplan.logical.relational.LogicalPlan;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.optimizer.Rule;
import org.apache.pig.newplan.optimizer.Transformer;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.impl.util.Utils;
public class MergeForEach extends Rule {
private OperatorSubPlan subPlan;
public MergeForEach(String name) {
super( name, false );
}
@Override
protected OperatorPlan buildPattern() {
// match each foreach.
LogicalPlan plan = new LogicalPlan();
LogicalRelationalOperator foreach1 = new LOForEach(plan);
plan.add( foreach1 );
return plan;
}
@Override
public Transformer getNewTransformer() {
return new MergeForEachTransformer();
}
public class MergeForEachTransformer extends Transformer {
@Override
public boolean check(OperatorPlan matched) throws FrontendException {
LOForEach foreach1 = (LOForEach)matched.getSources().get(0);
List<Operator> succs = currentPlan.getSuccessors( foreach1 );
if( succs == null || succs.size() != 1 || !( succs.get(0) instanceof LOForEach) )
return false;
LOForEach foreach2 = (LOForEach)succs.get(0);
// Check if the second foreach has only LOGenerate and LOInnerLoad
Iterator<Operator> it = foreach2.getInnerPlan().getOperators();
while( it.hasNext() ) {
Operator op = it.next();
if(!(op instanceof LOGenerate) && !(op instanceof LOInnerLoad))
return false;
}
// Check if the first foreach has flatten in its generate statement.
LOGenerate gen1 = (LOGenerate)foreach1.getInnerPlan().getSinks().get(0);
for (boolean flatten : gen1.getFlattenFlags()) {
if( flatten )
return false;
}
if (gen1.getUserDefinedSchema()!=null) {
for (LogicalSchema s : gen1.getUserDefinedSchema()) {
if (s!=null) {
return false;
}
}
}
// Check if non of the 1st foreach output is referred more than once in second foreach.
// Otherwise, we may do expression calculation more than once, defeat the benefit of this
// optimization
Set<Integer> inputs = new HashSet<Integer>();
boolean duplicateInputs = false;
for (Operator op : foreach2.getInnerPlan().getSources()) {
// If the source is not LOInnerLoad, then it must be LOGenerate. This happens when
// the 1st ForEach does not rely on any input of 2nd ForEach
if (op instanceof LOInnerLoad) {
LOInnerLoad innerLoad = (LOInnerLoad)op;
int input = innerLoad.getProjection().getColNum();
if (inputs.contains(input)) {
duplicateInputs = true;
break;
}
else
inputs.add(input);
if (innerLoad.getProjection().isRangeOrStarProject())
return false;
}
}
// Duplicate inputs in the case first foreach only containing LOInnerLoad and
// LOGenerate is allowed, and output plan is simple projection
if (duplicateInputs) {
Iterator<Operator> it1 = foreach1.getInnerPlan().getOperators();
while( it1.hasNext() ) {
Operator op = it1.next();
if(!(op instanceof LOGenerate) && !(op instanceof LOInnerLoad))
return false;
if (op instanceof LOGenerate) {
List<LogicalExpressionPlan> outputPlans = ((LOGenerate)op).getOutputPlans();
for (LogicalExpressionPlan outputPlan : outputPlans) {
Iterator<Operator> iter = outputPlan.getOperators();
while (iter.hasNext()) {
if (!(iter.next() instanceof ProjectExpression))
return false;
}
}
}
}
}
return true;
}
@Override
public OperatorPlan reportChanges() {
return subPlan;
}
// If op is LOInnerLoad, get a copy of it, otherwise, return op itself
private Operator getOperatorToMerge(Operator op, OperatorPlan newPlan, LOForEach newForEach) {
Operator opToMerge = op;
if (op instanceof LOInnerLoad) {
opToMerge = new LOInnerLoad(newPlan, newForEach, ((LOInnerLoad)op).getColNum());
} else {
opToMerge.setPlan(newPlan);
}
return opToMerge;
}
private Operator addBranchToPlan(LOGenerate gen, int branch, OperatorPlan newPlan, LOForEach newForEach) {
Operator opNextToGen;
Operator op = gen.getPlan().getPredecessors(gen).get(branch);
Operator opToMerge = getOperatorToMerge(op, newPlan, newForEach);
newPlan.add(opToMerge);
opNextToGen = opToMerge;
Operator pred;
if (gen.getPlan().getPredecessors(op)!=null)
pred = gen.getPlan().getPredecessors(op).get(0);
else
pred = null;
while (pred!=null) {
Operator predToMerge = getOperatorToMerge(pred, newPlan, newForEach);
newPlan.add(predToMerge);
newPlan.connect(predToMerge, op);
op = pred;
if (gen.getPlan().getPredecessors(pred)!=null)
pred = gen.getPlan().getPredecessors(pred).get(0);
else
pred = null;
}
return opNextToGen;
}
@Override
public void transform(OperatorPlan matched) throws FrontendException {
subPlan = new OperatorSubPlan(currentPlan);
LOForEach foreach1 = (LOForEach)matched.getSources().get(0);
LOGenerate gen1 = (LOGenerate)foreach1.getInnerPlan().getSinks().get(0);
LOForEach foreach2 = (LOForEach)currentPlan.getSuccessors(foreach1).get(0);
LOGenerate gen2 = (LOGenerate)foreach2.getInnerPlan().getSinks().get(0);
LOForEach newForEach = new LOForEach(currentPlan);
LogicalPlan newForEachInnerPlan = new LogicalPlan();
newForEach.setInnerPlan(newForEachInnerPlan);
newForEach.setAlias(foreach2.getAlias());
newForEach.setRequestedParallelism(foreach1.getRequestedParallelism());
List<LogicalExpressionPlan> newExpList = new ArrayList<LogicalExpressionPlan>();
LOGenerate newGen = new LOGenerate(newForEachInnerPlan, newExpList, gen2.getFlattenFlags());
newGen.setUserDefinedSchema(gen2.getUserDefinedSchema());
newForEachInnerPlan.add(newGen);
for (LogicalExpressionPlan exp2 : gen2.getOutputPlans()) {
LogicalExpressionPlan newExpPlan = new LogicalExpressionPlan();
LogicalExpressionPlan exp2Copy = exp2.deepCopy();
newExpPlan.merge(exp2Copy);
// Add expression plan in 2nd ForEach
List<Operator> exp2Sinks = new ArrayList<Operator>();
exp2Sinks.addAll(newExpPlan.getSinks());
for (Operator exp2Sink : exp2Sinks) {
if (exp2Sink instanceof ProjectExpression) {
// Find referred expression plan in 1st ForEach
ProjectExpression proj = (ProjectExpression)exp2Sink;
LOInnerLoad innerLoad = (LOInnerLoad)foreach2.getInnerPlan().getPredecessors(gen2).get(proj.getInputNum());
int exp1Pos = innerLoad.getProjection().getColNum();
LogicalExpressionPlan exp1 = gen1.getOutputPlans().get(exp1Pos);
LogicalExpressionPlan exp1Copy = exp1.deepCopy();
List<Operator> exp1Sources = newExpPlan.merge(exp1Copy);
// Copy expression plan to the new ForEach, connect to the expression plan of 2nd ForEach
Operator exp1Source = exp1Sources.get(0);
if (newExpPlan.getPredecessors(exp2Sink)!=null) {
Operator exp2NextToSink = newExpPlan.getPredecessors(exp2Sink).get(0);
Pair<Integer, Integer> pos = newExpPlan.disconnect(exp2NextToSink, exp2Sink);
newExpPlan.remove(exp2Sink);
newExpPlan.connect(exp2NextToSink, pos.first, exp1Source, 0);
}
else {
newExpPlan.remove(exp2Sink);
}
}
}
// Copy referred ForEach1 inner plan to new ForEach
List<Operator> exp1Sinks = newExpPlan.getSinks();
for (Operator exp1Sink : exp1Sinks) {
if (exp1Sink instanceof ProjectExpression) {
Operator opNextToGen = addBranchToPlan(gen1, ((ProjectExpression)exp1Sink).getInputNum(), newForEachInnerPlan, newForEach);
newForEachInnerPlan.connect(opNextToGen, newGen);
int input = newForEachInnerPlan.getPredecessors(newGen).indexOf(opNextToGen);
((ProjectExpression)exp1Sink).setInputNum(input);
}
}
newExpList.add(newExpPlan);
}
// Adjust attachedOp
for (LogicalExpressionPlan p : newGen.getOutputPlans()) {
Iterator<Operator> iter = p.getOperators();
while (iter.hasNext()) {
Operator op = iter.next();
if (op instanceof ProjectExpression) {
((ProjectExpression)op).setAttachedRelationalOp(newGen);
}
}
}
Iterator<Operator> iter = newForEach.getInnerPlan().getOperators();
while (iter.hasNext()) {
Operator op = iter.next();
if (op instanceof LOInnerLoad) {
((LOInnerLoad)op).getProjection().setAttachedRelationalOp(newForEach);
}
}
// remove foreach1, foreach2, add new foreach
// rebuild soft link
Collection<Operator> newSoftLinkPreds = Utils.mergeCollection(currentPlan.getSoftLinkPredecessors(foreach1),
currentPlan.getSoftLinkPredecessors(foreach2));
Collection<Operator> foreach1SoftLinkPred = null;
if (currentPlan.getSoftLinkPredecessors(foreach1)!=null) {
foreach1SoftLinkPred = new ArrayList<Operator>();
foreach1SoftLinkPred.addAll(currentPlan.getSoftLinkPredecessors(foreach1));
}
if (foreach1SoftLinkPred!=null) {
for (Operator softPred : foreach1SoftLinkPred) {
currentPlan.removeSoftLink(softPred, foreach1);
}
}
Collection<Operator> foreach2SoftLinkPred = null;
if (currentPlan.getSoftLinkPredecessors(foreach2)!=null) {
foreach2SoftLinkPred = new ArrayList<Operator>();
foreach2SoftLinkPred.addAll(currentPlan.getSoftLinkPredecessors(foreach2));
}
if (foreach2SoftLinkPred!=null) {
for (Operator softPred : foreach2SoftLinkPred) {
currentPlan.removeSoftLink(softPred, foreach2);
}
}
currentPlan.removeAndReconnect(foreach1);
currentPlan.replace(foreach2, newForEach);
if (newSoftLinkPreds!=null) {
for (Operator softPred : newSoftLinkPreds) {
currentPlan.createSoftLink(softPred, newForEach);
}
}
subPlan.add(newForEach);
}
}
}