/**
* 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.tez.plan;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.apache.pig.backend.hadoop.executionengine.tez.TezResourceManager;
import org.apache.pig.backend.hadoop.executionengine.tez.util.TezCompilerUtil;
import org.apache.pig.impl.plan.OperatorKey;
import org.apache.pig.impl.plan.OperatorPlan;
import org.apache.pig.impl.plan.PlanException;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.Pair;
/**
* A Plan used to create the plan of Tez operators which can be converted into
* the Job Control object. This is necessary to capture the dependencies among
* jobs.
*/
public class TezOperPlan extends OperatorPlan<TezOperator> {
private static final long serialVersionUID = 1L;
private Map<String, Path> extraResources = new HashMap<String, Path>();
private int estimatedTotalParallelism = -1;
public TezOperPlan() {
}
public int getEstimatedTotalParallelism() {
return estimatedTotalParallelism;
}
public void setEstimatedParallelism(int estimatedTotalParallelism) {
this.estimatedTotalParallelism = estimatedTotalParallelism;
}
@Override
public String toString() {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
PrintStream ps = new PrintStream(baos);
TezPrinter printer = new TezPrinter(ps, this);
printer.setVerbose(true);
try {
printer.visit();
} catch (VisitorException e) {
throw new RuntimeException(e);
}
return baos.toString();
}
// Add extra plan-specific local resources from the source FS
public void addExtraResource(URL url) throws IOException {
Path resourcePath = new Path(url.getFile());
String resourceName = resourcePath.getName();
if (!extraResources.containsKey(resourceName)) {
Path remoteFsPath;
try {
remoteFsPath = TezResourceManager.getInstance().addTezResource(url.toURI());
} catch (URISyntaxException e) {
throw new IOException(e);
}
extraResources.put(resourceName, remoteFsPath);
}
}
// Add extra plan-specific local resources already present in the remote FS
public void addExtraResource(String resourceName, Path remoteFsPath) throws IOException {
if (!extraResources.containsKey(resourceName)) {
TezResourceManager.getInstance().addTezResource(resourceName, remoteFsPath);
extraResources.put(resourceName, remoteFsPath);
}
}
// Get the plan-specific resources
public Map<String, LocalResource> getExtraResources() throws Exception {
// In a STREAM add the files specified in SHIP and CACHE
// as local resources for the plan.
TezPOStreamVisitor streamVisitor = new TezPOStreamVisitor(this);
streamVisitor.visit();
addShipResources(streamVisitor.getShipFiles());
addCacheResources(streamVisitor.getCacheFiles());
// In a UDF add the files specified by getCacheFiles()
// as local resources for the plan.
TezPOUserFuncVisitor udfVisitor = new TezPOUserFuncVisitor(this);
udfVisitor.visit();
addShipResources(udfVisitor.getShipFiles());
addCacheResources(udfVisitor.getCacheFiles());
return TezResourceManager.getInstance().getTezResources(extraResources.keySet());
}
// In the statement "SHIP('/home/foo')" we'll map the resource name foo to
// the file that has been copied to the staging directory in the remote FS.
private void addShipResources(Set<String> fileNames) throws IOException {
for (String fileName : fileNames) {
fileName = fileName.trim();
if (fileName.length() > 0) {
URL url = new File(fileName).toURI().toURL();
addExtraResource(url);
}
}
}
// In the statement "CACHE('/input/data.txt#alias.txt')" we'll map the
// resource name alias.txt to the actual resource path in the remote FS
// at '/input/data.txt'.
private void addCacheResources(Set<String> fileNames) throws Exception {
for (String fileName : fileNames) {
fileName = fileName.trim();
if (fileName.length() > 0) {
URI resourceURI = new URI(fileName);
String fragment = resourceURI.getFragment();
Path remoteFsPath = new Path(resourceURI.getPath());
String resourceName = (fragment != null && fragment.length() > 0) ? fragment : remoteFsPath.getName();
addExtraResource(resourceName, remoteFsPath);
}
}
}
@Override
public void remove(TezOperator op) {
// The remove method does not replace output and input keys in TezInput
// and TezOutput. That has to be handled separately.
for (OperatorKey opKey : op.outEdges.keySet()) {
getOperator(opKey).inEdges.remove(op.getOperatorKey());
}
for (OperatorKey opKey : op.inEdges.keySet()) {
getOperator(opKey).outEdges.remove(op.getOperatorKey());
}
super.remove(op);
}
@Override
public boolean disconnect(TezOperator from, TezOperator to) {
from.outEdges.remove(to.getOperatorKey());
to.inEdges.remove(from.getOperatorKey());
return super.disconnect(from, to);
}
/**
* Move everything below a given operator to the new operator plan. The specified operator will
* be moved and will be the root of the new operator plan
* @param root Operator to move everything under including the root operator
* @param newPlan new operator plan to move things into
* @throws PlanException
*/
public void moveTree(TezOperator root, TezOperPlan newPlan) throws PlanException {
List<TezOperator> list = new ArrayList<TezOperator>();
list.add(root);
int prevSize = 0;
int pos = 0;
while (list.size() > prevSize) {
prevSize = list.size();
TezOperator node = list.get(pos);
if (getSuccessors(node)!=null) {
for (TezOperator succ : getSuccessors(node)) {
if (!list.contains(succ)) {
list.add(succ);
}
}
}
if (getPredecessors(node)!=null) {
for (TezOperator pred : getPredecessors(node)) {
if (!list.contains(pred)) {
list.add(pred);
}
}
}
pos++;
}
for (TezOperator node: list) {
newPlan.add(node);
}
Set<Pair<TezOperator, TezOperator>> toReconnect = new HashSet<Pair<TezOperator, TezOperator>>();
for (TezOperator from : mFromEdges.keySet()) {
List<TezOperator> tos = mFromEdges.get(from);
for (TezOperator to : tos) {
if (list.contains(from) || list.contains(to)) {
toReconnect.add(new Pair<TezOperator, TezOperator>(from, to));
}
}
}
for (Pair<TezOperator, TezOperator> pair : toReconnect) {
if (list.contains(pair.first) && list.contains(pair.second)) {
// Need to reconnect in newPlan
TezEdgeDescriptor edge = pair.second.inEdges.get(pair.first.getOperatorKey());
TezCompilerUtil.connect(newPlan, pair.first, pair.second, edge);
}
}
for (TezOperator node : list) {
// Simply remove from plan, don't deal with inEdges/outEdges
super.remove(node);
}
}
}