/*
* 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.visitor;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import org.apache.pig.PigException;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.newplan.DependencyOrderWalker;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.PlanWalker;
import org.apache.pig.newplan.ReverseDependencyOrderWalker;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.LogicalExpressionVisitor;
import org.apache.pig.newplan.logical.expression.ProjectExpression;
import org.apache.pig.newplan.logical.expression.UserFuncExpression;
import org.apache.pig.newplan.logical.optimizer.AllExpressionVisitor;
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;
/**
* Expand project-star or project-range when used as udf argument.
* This is different from {@link ProjectStarExpander} because in those
* cases, the project star gets expanded as new {@link LogicalExpressionPlan}.
* In case of project-star or project-range within udf, it should get expanded
* only as multiple inputs to this udf, no addtional {@link LogicalExpressionPlan}s
* are created.
* The expansion happens only if input schema is not null
*/
public class ProjStarInUdfExpander extends AllExpressionVisitor {
public ProjStarInUdfExpander(OperatorPlan plan) throws FrontendException {
super( plan, new DependencyOrderWalker( plan ) );
}
@Override
protected LogicalExpressionVisitor getVisitor(final LogicalExpressionPlan exprPlan)
throws FrontendException {
//This handles the expansion udf in all operators other than foreach
return new ProjExpanderForNonForeach(exprPlan);
}
/*
* LOForeach needs special handling because LOInnerLoad's inner ProjectExpression
* is the one that gets expanded
*/
@Override
public void visit(LOForEach foreach) throws FrontendException{
LogicalPlan innerPlan = foreach.getInnerPlan();
//visit the inner plan first
PlanWalker newWalker = currentWalker.spawnChildWalker(innerPlan);
pushWalker(newWalker);
currentWalker.walk(this);
popWalker();
//get the LOGenerate
List<Operator> feOutputs = innerPlan.getSinks();
LOGenerate gen = null;
for( Operator op : feOutputs){
if(op instanceof LOGenerate){
if(gen != null){
String msg = "Expected single LOGenerate output in innerplan of foreach";
throw new VisitorException(foreach,
msg,
2266,
PigException.BUG
);
}
gen = (LOGenerate) op;
}
}
List<Operator> loGenPreds = innerPlan.getPredecessors(gen);
if(loGenPreds == null){
// there are no LOInnerLoads , must be working on just constants
// no project-star expansion to be done
return;
}
//get mapping of LOGenerate predecessor current position to object
Map<Integer, LogicalRelationalOperator> oldPos2Rel =
new HashMap<Integer, LogicalRelationalOperator>();
for(int i=0; i<loGenPreds.size(); i++){
oldPos2Rel.put(i, (LogicalRelationalOperator) loGenPreds.get(i));
}
//store mapping between the projection in inner plans of
// of LOGenerate to the input relation object
Map<ProjectExpression, LogicalRelationalOperator> proj2InpRel =
new HashMap<ProjectExpression, LogicalRelationalOperator>();
List<LOInnerLoad> expandedInLoads = new ArrayList<LOInnerLoad>();
//visit each expression plan, and expand the projects in the udf
for( OperatorPlan plan : gen.getOutputPlans()){
ProjExpanderForForeach projExpander = new ProjExpanderForForeach(
plan,
gen,
oldPos2Rel,
proj2InpRel,
foreach,
expandedInLoads
);
projExpander.visit();
}
//remove the LOInnerLoads that have been expanded
for(LOInnerLoad inLoad : expandedInLoads){
innerPlan.disconnect(inLoad, gen);
innerPlan.remove(inLoad);
}
//reset the input relation position in the projects
//get mapping of LoGenerate input relation to current position
Map<LogicalRelationalOperator, Integer> rel2pos = new HashMap<LogicalRelationalOperator, Integer>();
List<Operator> newGenPreds = innerPlan.getPredecessors(gen);
int numNewGenPreds = 0;
if(newGenPreds != null)
numNewGenPreds = newGenPreds.size();
for(int i=0; i<numNewGenPreds; i++){
rel2pos.put((LogicalRelationalOperator) newGenPreds.get(i),i);
}
//correct the input num for projects
for(Entry<ProjectExpression, LogicalRelationalOperator> projAndInp : proj2InpRel.entrySet()){
ProjectExpression proj = projAndInp.getKey();
LogicalRelationalOperator rel = projAndInp.getValue();
proj.setInputNum(rel2pos.get(rel));
}
}
@Override
public void visit(LOGenerate gen) throws FrontendException{
}
}
class ProjExpanderForForeach extends LogicalExpressionVisitor{
private LOGenerate loGen;
private LogicalPlan innerRelPlan;
private Map<Integer, LogicalRelationalOperator> oldPos2Rel;
private Map<ProjectExpression, LogicalRelationalOperator> proj2InpRel;
private LOForEach foreach;
private List<LOInnerLoad> expandedInLoads;
protected ProjExpanderForForeach(
OperatorPlan p,
LOGenerate loGen,
Map<Integer, LogicalRelationalOperator> oldPos2Rel,
Map<ProjectExpression, LogicalRelationalOperator> proj2InpRel,
LOForEach foreach,
List<LOInnerLoad> expandedInLoads
)
throws FrontendException {
super(p, new ReverseDependencyOrderWalker(p));
this.loGen = loGen;
this.innerRelPlan = (LogicalPlan) loGen.getPlan();
this.oldPos2Rel = oldPos2Rel;
this.proj2InpRel = proj2InpRel;
this.foreach = foreach;
this.expandedInLoads = expandedInLoads;
}
@Override
public void visit(UserFuncExpression func) throws FrontendException{
if(plan.getSuccessors(func) == null){
// no args for the udf, so nothing to do
return;
}
List<Operator> inputs = new ArrayList<Operator>(plan.getSuccessors(func));
// expandedProjectStars will be removed from plan
List<Operator> expandedProjectStars = new ArrayList<Operator>();
// new projects to be added to the plan
List<Operator> newExpandedProjects = new ArrayList<Operator>();
//new set of inputs
List<Operator> newInputs = new ArrayList<Operator>();
for(Operator inp : inputs){
if(inp instanceof ProjectExpression && ((ProjectExpression)inp).isRangeOrStarProject()
&& oldPos2Rel.get(((ProjectExpression)inp).getInputNum()) instanceof LOInnerLoad
){
//under foreach the ProjectExpression is always a project-star,
// need to check if the input relation is a LOInnerLoad
// containing project-star/range
LOInnerLoad inLoad =
(LOInnerLoad)oldPos2Rel.get(((ProjectExpression)inp).getInputNum());
ProjectExpression innerProj = inLoad.getProjection();
if(!innerProj.isRangeOrStarProject()){
newInputs.add(inp);
continue;
}
//try expanding the project-star/range
List<Operator> expandedOps = expandProjectStar(innerProj);
if(expandedOps != null){
expandedProjectStars.add(inp);//to remove
expandedInLoads.add(inLoad);//to remove
newInputs.addAll(expandedOps);
newExpandedProjects.addAll(expandedOps);//to add
}else {
newInputs.add(inp);
}
}else{
newInputs.add(inp);
}
}
//make changes to the plan if there is a project that was expanded
if(expandedProjectStars.size() > 0){
//disconnect old inputs
for(Operator inp : inputs){
plan.disconnect(func, inp);
}
//remove expanded projects
for(Operator op : expandedProjectStars){
plan.remove(op);
proj2InpRel.remove(op);
}
//add new projects
for(Operator op : newExpandedProjects){
plan.add(op);
}
//connect new inputs
for(Operator newInp : newInputs){
plan.connect(func, newInp);
}
}
}
@Override
public void visit(ProjectExpression proj){
//add project to LOInnerLoad mapping so that the input number can be
//corrected later
proj2InpRel.put(proj, oldPos2Rel.get(proj.getInputNum()));
}
private List<Operator> expandProjectStar(ProjectExpression proj)
throws FrontendException {
Pair<Integer, Integer> firstLastCols =
ProjectStarExpanderUtil.getProjectStartEndCols((LogicalExpressionPlan)plan, proj);
if(firstLastCols == null){
//no expansion happening now
return null;
}
//expand from firstProjCol to lastProjCol
int firstProjCol = firstLastCols.first;
int lastProjCol = firstLastCols.second;
List<Operator> newProjects = new ArrayList<Operator>();
for(int i = firstProjCol; i <= lastProjCol; i++){
LOInnerLoad newILoad = new LOInnerLoad(innerRelPlan, foreach, i);
innerRelPlan.add(newILoad);
innerRelPlan.connect(newILoad, loGen);
ProjectExpression newProj = new ProjectExpression(plan, -2, -1, loGen) ;
proj2InpRel.put(newProj, newILoad);
newProjects.add(newProj);
}
return newProjects;
}
}
class ProjExpanderForNonForeach extends LogicalExpressionVisitor{
protected ProjExpanderForNonForeach(OperatorPlan p)
throws FrontendException {
super(p, new ReverseDependencyOrderWalker(p));
}
@Override
public void visit(UserFuncExpression func) throws FrontendException {
if(plan.getSuccessors(func) == null){
//udf without args, nothing to do
return;
}
List<Operator> inputs = new ArrayList<Operator>(plan.getSuccessors(func));
// expandedProjectStars will be removed from plan
List<Operator> expandedProjectStars = new ArrayList<Operator>();
// new projects to be added to the plan
List<Operator> newExpandedProjects = new ArrayList<Operator>();
//new set of inputs
List<Operator> newInputs = new ArrayList<Operator>();
for(Operator inp : inputs){
if(inp instanceof ProjectExpression && ((ProjectExpression)inp).isRangeOrStarProject() ){
//try expanding the project-star/range
List<Operator> expandedOps = expandProjectStar((ProjectExpression)inp);
if(expandedOps != null){
expandedProjectStars.add(inp);//to remove
newInputs.addAll(expandedOps);
newExpandedProjects.addAll(expandedOps);//to add
}else {
newInputs.add(inp);
}
}else{
newInputs.add(inp);
}
}
//make changes to the plan if there is a project that was expanded
if(expandedProjectStars.size() > 0){
//disconnect old inputs
for(Operator inp : inputs){
plan.disconnect(func, inp);
}
//remove expanded projects
for(Operator op : expandedProjectStars){
plan.remove(op);
}
//add new projects
for(Operator op : newExpandedProjects){
plan.add(op);
}
//connect new inputs
for(Operator newInp : newInputs){
plan.connect(func, newInp);
}
}
}
private List<Operator> expandProjectStar(
ProjectExpression proj) throws FrontendException {
Pair<Integer, Integer> firstLastCols =
ProjectStarExpanderUtil.getProjectStartEndCols((LogicalExpressionPlan)plan, proj);
if(firstLastCols == null){
//no expansion happening now
return null;
}
//expand from firstProjCol to lastProjCol
int firstProjCol = firstLastCols.first;
int lastProjCol = firstLastCols.second;
List<Operator> newProjects = new ArrayList<Operator>();
LogicalRelationalOperator relOp = proj.getAttachedRelationalOp();
for(int i = firstProjCol; i <= lastProjCol; i++){
newProjects.add(new ProjectExpression(plan, proj.getInputNum(), i, relOp));
}
return newProjects;
}
};