/*
* 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.backend.hadoop.executionengine.mapReduceLayer;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.pig.FuncSpec;
import org.apache.pig.backend.hadoop.datastorage.ConfigurationUtil;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROpPlanVisitor;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROperPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.ConstantExpression;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserFunc;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhyPlanVisitor;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POForEach;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLoad;
import org.apache.pig.impl.io.FileSpec;
import org.apache.pig.impl.plan.DepthFirstWalker;
import org.apache.pig.impl.plan.PlanException;
import org.apache.pig.impl.plan.PlanWalker;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.Utils;
import org.apache.pig.impl.PigContext;
/**
* A visitor to optimize plans that have a sample job that immediately follows a
* load/store only MR job. These kinds of plans are generated for order bys, and
* will soon be generated for joins that need to sample their data first. These
* can be changed so that the RandomSampleLoader subsumes the loader used in the
* first job and then removes the first job.
*/
public class SampleOptimizer extends MROpPlanVisitor {
private static final Log log = LogFactory.getLog(SampleOptimizer.class);
private PigContext pigContext;
public SampleOptimizer(MROperPlan plan, PigContext pigContext) {
super(plan, new DepthFirstWalker<MapReduceOper, MROperPlan>(plan));
this.pigContext = pigContext;
}
private List<MapReduceOper> opsToRemove = new ArrayList<MapReduceOper>();
@Override
public void visit() throws VisitorException {
super.visit();
// remove operators marked for removal
for (MapReduceOper op : opsToRemove)
this.mPlan.remove(op);
}
@Override
public void visitMROp(MapReduceOper mr) throws VisitorException {
// See if this is a sampling job.
List<PhysicalOperator> pos = mr.mapPlan.getRoots();
if (pos == null || pos.size() == 0) {
log.debug("Map of operator empty");
return;
}
PhysicalOperator po = pos.get(0);
if (!(po instanceof POLoad)) {
log.debug("Root operator of map is not load.");
return; // Huh?
}
POLoad load = (POLoad)po;
String loadFunc = load.getLFile().getFuncName();
String loadFile = load.getLFile().getFileName();
if (!("org.apache.pig.impl.builtin.RandomSampleLoader".equals(loadFunc)) && !("org.apache.pig.impl.builtin.PoissonSampleLoader".equals(loadFunc))) {
log.debug("Not a sampling job.");
return;
}
if (loadFile == null) {
log.debug("No load file");
return;
}
// Get this job's predecessor. There should be exactly one.;
List<MapReduceOper> preds = mPlan.getPredecessors(mr);
if (preds.size() != 1) {
log.debug("Too many predecessors to sampling job.");
return;
}
MapReduceOper pred = preds.get(0);
// The predecessor should be a root.
List<MapReduceOper> predPreds = mPlan.getPredecessors(pred);
if (predPreds != null && predPreds.size() > 0) {
log.debug("Predecessor should be a root of the plan");
return;
}
// The predecessor should have just a load and store in the map, and nothing
// in the combine or reduce.
if ( !(pred.reducePlan.isEmpty() && pred.combinePlan.isEmpty())) {
log.debug("Predecessor has a combine or reduce plan");
return;
}
// The MR job should have one successor.
List<MapReduceOper> succs = mPlan.getSuccessors(mr);
if (succs.size() != 1) {
log.debug("Job has more than one successor.");
return;
}
MapReduceOper succ = succs.get(0);
if (pred.mapPlan == null || pred.mapPlan.size() != 2) {
log.debug("Predecessor has more than just load+store in the map");
return;
}
List<PhysicalOperator> loads = pred.mapPlan.getRoots();
if (loads.size() != 1) {
log.debug("Predecessor plan has more than one root.");
return;
}
PhysicalOperator r = loads.get(0);
if (!(r instanceof POLoad)) { // Huh?
log.debug("Predecessor's map plan root is not a load.");
return;
}
POLoad predLoad = (POLoad)r;
// Find the load the correlates with the file the sampler is loading, and
// check that it is using the temp file storage format.
if (succ.mapPlan == null) { // Huh?
log.debug("Successor has no map plan.");
return;
}
loads = succ.mapPlan.getRoots();
POLoad succLoad = null;
for (PhysicalOperator root : loads) {
if (!(root instanceof POLoad)) { // Huh?
log.debug("Successor's roots are not loads");
return;
}
POLoad sl = (POLoad)root;
if (loadFile.equals(sl.getLFile().getFileName()) &&
Utils.getTmpFileCompressorName(pigContext).equals(sl.getLFile().getFuncName())) {
succLoad = sl;
break;
}
}
if (succLoad == null) {
log.debug("Could not find load that matched file we are sampling.");
return;
}
// Okay, we're on.
// First, replace our RandomSampleLoader with a RandomSampleLoader that uses
// the load function from our predecessor.
String[] rslargs = new String[2];
FileSpec predFs = predLoad.getLFile();
// First argument is FuncSpec of loader function to subsume, this we want to set for
// ourselves.
rslargs[0] = predFs.getFuncSpec().toString();
// Add the loader's funcspec to the list of udf's associated with this mr operator
mr.UDFs.add(rslargs[0]);
// Second argument is the number of samples per block, read this from the original.
rslargs[1] = load.getLFile().getFuncSpec().getCtorArgs()[1];
FileSpec fs = new FileSpec(predFs.getFileName(),new FuncSpec(loadFunc, rslargs));
POLoad newLoad = new POLoad(load.getOperatorKey(),load.getRequestedParallelism(), fs);
newLoad.setSignature(predLoad.getSignature());
newLoad.setLimit(predLoad.getLimit());
try {
mr.mapPlan.replace(load, newLoad);
// check if it has PartitionSkewedKeys
List<PhysicalOperator> ls = mr.reducePlan.getLeaves();
for(PhysicalOperator op: ls) {
scan(mr, op, fs.getFileName());
}
} catch (PlanException e) {
throw new VisitorException(e);
}
// Second, replace the loader in our successor with whatever the originally used loader was.
fs = new FileSpec(predFs.getFileName(), predFs.getFuncSpec());
newLoad = new POLoad(succLoad.getOperatorKey(), succLoad.getRequestedParallelism(), fs);
newLoad.setSignature(predLoad.getSignature());
try {
succ.mapPlan.replace(succLoad, newLoad);
// Add the loader's funcspec to the list of udf's associated with this mr operator
succ.UDFs.add(newLoad.getLFile().getFuncSpec().toString());
} catch (PlanException e) {
throw new VisitorException(e);
}
// Cannot delete the pred right now, because we are still traversing the graph. So, mark the pred and remove it from the
// plan once the visit by this optimizer is complete.
opsToRemove.add(pred);
}
// search for PartionSkewedKeys and update input file name
// it is always under a POForEach operator in reduce plan
private void scan(MapReduceOper mr, PhysicalOperator op, String fileName) {
if (op instanceof POUserFunc) {
if (((POUserFunc)op).getFuncSpec().getClassName().equals(
"org.apache.pig.impl.builtin.PartitionSkewedKeys")) {
String[] ctorArgs = ((POUserFunc)op).getFuncSpec().getCtorArgs();
ctorArgs[2] = fileName;
return;
}
}else if (op instanceof POForEach) {
List<PhysicalPlan> pl = ((POForEach)op).getInputPlans();
for(PhysicalPlan plan: pl) {
List<PhysicalOperator> list = plan.getLeaves();
for (PhysicalOperator pp: list) {
scan(mr, pp, fileName);
}
}
}else{
List<PhysicalOperator> preds = mr.reducePlan.getPredecessors(op);
if (preds == null) {
return;
}
for(PhysicalOperator p: preds) {
scan(mr, p, fileName);
}
}
}
}