/*
* 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.test;
import static org.junit.Assert.*;
import java.io.IOException;
import java.util.*;
import org.apache.pig.ExecType;
import org.apache.pig.PigServer;
import org.apache.pig.newplan.logical.optimizer.LogicalPlanOptimizer;
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.LogicalPlan;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.logical.rules.AddForEach;
import org.apache.pig.newplan.logical.rules.LoadTypeCastInserter;
import org.apache.pig.newplan.logical.rules.MergeForEach;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.optimizer.PlanOptimizer;
import org.apache.pig.newplan.optimizer.Rule;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.impl.PigContext;
import junit.framework.Assert;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
public class TestMergeForEachOptimization {
LogicalPlan plan = null;
PigContext pc = new PigContext( ExecType.LOCAL, new Properties() );
PigServer pigServer = null;
@Before
public void setup() throws ExecException {
pigServer = new PigServer( pc );
}
@After
public void tearDown() {
}
/**
* Basic test case. Two simple FOREACH statements can be merged to one.
* @throws Exception
*/
@Test
public void testSimple() throws Exception {
String query = "A = load 'file.txt' as (a, b, c);" +
"B = foreach A generate a+b, c-b;" +
"C = foreach B generate $0+5, $1;" +
"store C into 'empty';";
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
int outputExprCount1 = getOutputExprCount( newLogicalPlan );
LOForEach foreach1 = getForEachOperator( newLogicalPlan );
Assert.assertTrue( foreach1.getAlias().equals( "C" ) );
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
Assert.assertEquals( 1, forEachCount1 - forEachCount2 );
int outputExprCount2 = getOutputExprCount( newLogicalPlan );
Assert.assertTrue( outputExprCount1 == outputExprCount2 );
LOForEach foreach2 = getForEachOperator( newLogicalPlan );
Assert.assertTrue( foreach2.getAlias().equals( "C" ) );
}
/**
* Test more complex case where the first for each in the script has inner plan.
* @throws Exception
*/
@Test
public void testComplex() throws Exception {
String query = "A = load 'file.txt' as (a:int, b, c:bag{t:tuple(c0:int,c1:int)});" +
"B = foreach A { S = ORDER c BY $0; generate $0, COUNT(S), SIZE(S); };" +
"C = foreach B generate $2+5 as x, $0-$1/2 as y;" + "store C into 'empty';" ;
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
int outputExprCount1 = getOutputExprCount( newLogicalPlan );
LOForEach foreach1 = getForEachOperator( newLogicalPlan );
Assert.assertTrue( foreach1.getAlias().equals( "C" ) );
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
// The number of FOREACHes didn't change because one is genereated because of type cast and
// one is reduced because of the merge.
Assert.assertEquals( 0, forEachCount1 - forEachCount2 );
int outputExprCount2 = getOutputExprCount( newLogicalPlan );
Assert.assertTrue( outputExprCount1 == outputExprCount2 );
LOForEach foreach2 = getForEachOperator( newLogicalPlan );
Assert.assertTrue( foreach2.getAlias().equals( "C" ) );
LogicalSchema newSchema = foreach2.getSchema();
Assert.assertTrue(newSchema.getField(0).alias.equals("x"));
Assert.assertTrue(newSchema.getField(1).alias.equals("y"));
}
/**
* One output of first foreach was referred more than once in the second foreach
* @throws Exception
*/
@Test
public void testDuplicateInputs() throws Exception {
String query = "A = load 'file.txt' as (a0:int, a1:double);" +
"A1 = foreach A generate (int)a0 as a0, (double)a1 as a1;" +
"B = group A1 all;" +
"C = foreach B generate A1;" +
"D = foreach C generate SUM(A1.a0), AVG(A1.a1);" + "store D into 'empty';" ;
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
Operator store = newLogicalPlan.getSinks().get(0);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
LOForEach foreach1 = (LOForEach)newLogicalPlan.getPredecessors(store).get(0);
Assert.assertTrue( foreach1.getAlias().equals( "D" ) );
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
// The number of FOREACHes didn't change because one is genereated because of type cast and
// one is reduced because of the merge.
Assert.assertEquals( 1, forEachCount1 - forEachCount2 );
LOForEach foreach2 = (LOForEach)newLogicalPlan.getPredecessors(store).get(0);
Assert.assertTrue( foreach2.getAlias().equals( "D" ) );
}
/**
* Not all consecutive FOREACHes can be merged. In this case, the second FOREACH statment
* has inner plan, which cannot be merged with one before it.
* @throws Exception
*/
@Test
public void testNegative1() throws Exception {
String query = "A = LOAD 'file.txt' as (a, b, c, d:bag{t:tuple(c0:int,c1:int)});" +
"B = FOREACH A GENERATE a+5 AS u, b-c/2 AS v, d AS w;" +
"C = FOREACH B { S = ORDER w BY $0; GENERATE $0 as x, COUNT(S) as y; };" +
"store C into 'empty';";
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
// Actually MergeForEach optimization is happening here. A new foreach will be inserted after A because
// of typ casting. The inserted one and the one in B can be merged due to this optimization. However,
// the plan cannot be further optimized because C has inner plan.
Assert.assertEquals( forEachCount1, forEachCount2 );
}
/**
* MergeForEach Optimization is off if the first statement has a FLATTEN operator.
* @throws Exception
*/
@Test
public void testNegative2() throws Exception {
String query = "A = LOAD 'file.txt' as (a, b, c);" +
"B = FOREACH A GENERATE FLATTEN(a), b, c;" +
"C = FOREACH B GENERATE $0, $1+$2;" + "store C into 'empty';" ;
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
Assert.assertEquals( 2, forEachCount1 );
Assert.assertEquals( 2, forEachCount2 );
}
/**
* Ensure that join input order does not get reversed (PIG-1672)
* @throws Exception
*/
@Test
public void testJoinInputOrder() throws Exception {
String query = "l1 = load 'y' as (a);" +
"l2 = load 'z' as (a1,b1,c1,d1);" +
"f1 = foreach l2 generate a1, b1, c1, d1;" +
"f2 = foreach f1 generate a1, b1, c1;" +
"j1 = join f2 by a1, l1 by a using 'replicated';" + "store j1 into 'empty';" ;
LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
int forEachCount1 = getForEachOperatorCount( newLogicalPlan );
List<Operator> loads = newLogicalPlan.getSources();
Operator l2 = null;
for (Operator load : loads) {
if (((LogicalRelationalOperator)load).getAlias().equals("l2"))
l2 = load;
}
LOForEach foreachL2 = (LOForEach)newLogicalPlan.getSuccessors(l2).get(0);
foreachL2 = (LOForEach)newLogicalPlan.getSuccessors(foreachL2).get(0);
int outputExprCount1 = ((LOGenerate)foreachL2.getInnerPlan().getSinks().get(0)).getOutputPlans().size();
PlanOptimizer optimizer = new MyPlanOptimizer( newLogicalPlan, 3 );
optimizer.optimize();
int forEachCount2 = getForEachOperatorCount( newLogicalPlan );
Assert.assertEquals( 0, forEachCount1 - forEachCount2 );
loads = newLogicalPlan.getSources();
l2 = null;
for (Operator load : loads) {
if (((LogicalRelationalOperator)load).getAlias().equals("l2"))
l2 = load;
}
foreachL2 = (LOForEach)newLogicalPlan.getSuccessors(l2).get(0);
int outputExprCount2 = ((LOGenerate)foreachL2.getInnerPlan().getSinks().get(0)).getOutputPlans().size();
Assert.assertTrue( outputExprCount1 == outputExprCount2 );
Assert.assertTrue( foreachL2.getAlias().equals( "f2" ) );
LOJoin join = (LOJoin)getOperator(newLogicalPlan, LOJoin.class);
LogicalRelationalOperator leftInp =
(LogicalRelationalOperator)newLogicalPlan.getPredecessors(join).get(0);
assertEquals("join child left", leftInp.getAlias(), "f2");
LogicalRelationalOperator rightInp =
(LogicalRelationalOperator)newLogicalPlan.getPredecessors(join).get(1);
assertEquals("join child right", rightInp.getAlias(), "l1");
}
private int getForEachOperatorCount(LogicalPlan plan) {
Iterator<Operator> ops = plan.getOperators();
int count = 0;
while( ops.hasNext() ) {
Operator op = ops.next();
if( op instanceof LOForEach )
count++;
}
return count;
}
private int getOutputExprCount(LogicalPlan plan) throws IOException {
LOForEach foreach = getForEachOperator( plan );
LogicalPlan inner = foreach.getInnerPlan();
List<Operator> ops = inner.getSinks();
LOGenerate gen = (LOGenerate)ops.get( 0 );
return gen.getOutputPlans().size();
}
private LOForEach getForEachOperator(LogicalPlan plan) throws IOException {
Iterator<Operator> ops = plan.getOperators();
while( ops.hasNext() ) {
Operator op = ops.next();
if( op instanceof LOForEach ) {
LOForEach foreach = (LOForEach)op;
Operator succ = plan.getSuccessors( foreach ).get( 0 );
if( !(succ instanceof LOForEach ) )
return foreach;
}
}
return null;
}
/**
* returns first operator that is an instance of given class c
* @param plan
* @param c
* @return instance of class c if any, or null
* @throws IOException
*/
private Operator getOperator(LogicalPlan plan, Class<? extends Operator> c) throws IOException {
Iterator<Operator> ops = plan.getOperators();
while( ops.hasNext() ) {
Operator op = ops.next();
if( op.getClass().equals(c)) {
return op;
}
}
return null;
}
public class MyPlanOptimizer extends LogicalPlanOptimizer {
protected MyPlanOptimizer(OperatorPlan p, int iterations) {
super(p, iterations, new HashSet<String>());
}
protected List<Set<Rule>> buildRuleSets() {
List<Set<Rule>> ls = new ArrayList<Set<Rule>>();
Set<Rule> s = new HashSet<Rule>();
// add split filter rule
Rule r = new LoadTypeCastInserter( "TypeCastInserter" );
s.add(r);
ls.add(s);
// Split Set
// This set of rules does splitting of operators only.
// It does not move operators
s = new HashSet<Rule>();
r = new AddForEach( "AddForEach" );
s.add(r);
ls.add(s);
s = new HashSet<Rule>();
r = new MergeForEach("MergeForEach");
s.add(r);
ls.add(s);
return ls;
}
}
}