test/org/apache/pig/test/TestNewPlanColumnPrune.java - pig - Git at Google

 /*
  * 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.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;

 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;

 import org.apache.pig.PigServer;
 import org.apache.pig.impl.PigContext;
 import org.apache.pig.newplan.Operator;
 import org.apache.pig.newplan.OperatorPlan;
 import org.apache.pig.newplan.logical.optimizer.LogicalPlanOptimizer;
 import org.apache.pig.newplan.logical.relational.LOLoad;
 import org.apache.pig.newplan.logical.relational.LogicalPlan;
 import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
 import org.apache.pig.newplan.logical.rules.AddForEach;
 import org.apache.pig.newplan.logical.rules.ColumnMapKeyPrune;
 import org.apache.pig.newplan.logical.rules.MapKeysPruneHelper;
 import org.apache.pig.newplan.optimizer.PlanOptimizer;
 import org.apache.pig.newplan.optimizer.Rule;
 import org.junit.Before;
 import org.junit.Test;

 public class TestNewPlanColumnPrune {
     LogicalPlan plan = null;
     PigContext pc;

     @Before
     public void setUp() throws Exception {
         pc = new PigContext(Util.getLocalTestMode(), new Properties());
     }
     private LogicalPlan buildPlan(String query) throws Exception{
         PigServer pigServer = new PigServer( pc );
         return Util.buildLp(pigServer, query);
     }

     @Test
     public void testNoPrune() throws Exception  {
         // no foreach
         String query = "a = load 'd.txt' as (id, v1, v2);" +
         "b = filter a by v1==NULL;" +
         "store b into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query  = "a = load 'd.txt' as (id, v1, v2);" +
         "b = filter a by v1==NULL;"  +
         "store b into 'empty';";
         LogicalPlan expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));

         // no schema
         query = "a = load 'd.txt';" +
         "b = foreach a generate $0, $1;" +
         "store b into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a = load 'd.txt';"+
         "b = foreach a generate $0, $1;"+
         "store b into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));
     }

     @Test
     public void testPrune() throws Exception  {
         // only foreach
         String query = "a = load 'd.txt' as (id, v1, v2);" +
         "b = foreach a generate id;"+
         "store b into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a = load 'd.txt' as (id);" +
         "b = foreach a generate id;"+
         "store b into 'empty';";
         LogicalPlan expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));

         // with filter
         query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);"+
         "b = filter a by v1 != NULL AND (v2+v3)<100;"+
         "c = foreach b generate id;"+
         "store c into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a = load 'd.txt' as (id, v1, v3, v2);" +
         "b = filter a by v1 != NULL AND (v2+v3)<100;" +
         "c = foreach b generate id;" +
         "store c into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

         // with 2 foreach
         query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v2, v5, v4;" +
         "c = foreach b generate v5, v4;" +
         "store c into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a = load 'd.txt' as (v5, v4);" +
         "b = foreach a generate v5, v4;" +
         "c = foreach b generate v5, v4;" +
         "store c into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

         // with 2 foreach
         query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate id, v1, v5, v3, v4;" +
         "c = foreach b generate v5, v4;" +
         "store c into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a = load 'd.txt' as (v5, v4);" +
         "b = foreach a generate v5, v4;" +
         "c = foreach b generate v5, v4;" +
         "store c into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

         // with 2 foreach and filter in between
         query = "a =load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v2, v5, v4;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5, v4;" +
         "store d into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (v5, v4, v2);" +
         "b = foreach a generate v2, v5, v4;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5, v4;" +
         "store d into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

         // with 2 foreach after join
         query = "a =load 'd.txt' as (id, v1, v2, v3);" +
         "b = load 'c.txt' as (id, v4, v5, v6);" +
         "c = join a by id, b by id;" +
         "d = foreach c generate a::id, v5, v3, v4;" +
         "store d into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, v3);" +
         "b = load 'c.txt' as (id, v4, v5);" +
         "c = join a by id, b by id;" +
         "d = foreach c generate a::id, v5, v3, v4;" +
         "store d into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

         // with BinStorage, insert foreach after load
         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "c = filter a by v2 != NULL;" +
         "d = foreach c generate v5, v4;" +
         "store d into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v5, v4, v2;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5, v4;" +
         "store d into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

        // with BinStorage, not to insert foreach after load if there is already one
         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v5, v4, v2;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5;" +
         "store d into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v5, v2;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5;" +
         "store d into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));

        // with BinStorage, not to insert foreach after load if there is already one
         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v5, v4, v2, 10;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5;" +
         "store d into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
         "b = foreach a generate v5, v2, 10;" +
         "c = filter b by v2 != NULL;" +
         "d = foreach c generate v5;" +
         "store d into 'empty';";
         expected = buildPlan(query);
         assertTrue(expected.isEqual(newLogicalPlan));
     }

     @Test
     @SuppressWarnings("unchecked")
     public void testPruneWithMapKey() throws Exception {
          // only foreach
         String query = "a =load 'd.txt' as (id, v1, m:map[]);" +
         "b = foreach a generate id, m#'path';" +
         "store b into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, m:map[]);" +
         "b = foreach a generate id, m#'path';" +
         "store b into 'empty';";
         LogicalPlan expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));

         LOLoad op = (LOLoad)newLogicalPlan.getSources().get(0);
         Map<Integer,Set<String>> annotation =
                 (Map<Integer, Set<String>>) op.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
         assertEquals(1, annotation.size());
         Set<String> s = new HashSet<String>();
         s.add("path");
         assertEquals(annotation.get(2), s);

         // foreach with join
         query = "a =load 'd.txt' as (id, v1, m:map[]);" +
         "b = load 'd.txt' as (id, v1, m:map[]);" +
         "c = join a by id, b by id;" +
         "d = filter c by a::m#'path' != NULL;" +
         "e = foreach d generate a::id, b::id, b::m#'path', a::m;" +
         "store e into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, m:map[]);" +
         "b = load 'd.txt' as (id, m:map[]);" +
         "c = join a by id, b by id;" +
         "d = filter c by a::m#'path' != NULL;" +
         "e = foreach d generate a::id, b::id, b::m#'path', a::m;" +
         "store e into 'empty';";
         expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));

         List<Operator> ll = newLogicalPlan.getSources();
         assertEquals(2, ll.size());
         LOLoad loada = null;
         LOLoad loadb = null;
         for(Operator opp: ll) {
             if (((LogicalRelationalOperator)opp).getAlias().equals("a")) {
                 loada = (LOLoad)opp;
                 continue;
             }

             if (((LogicalRelationalOperator)opp).getAlias().equals("b")) {
                 loadb = (LOLoad)opp;
                 continue;
             }
         }

         annotation =
                 (Map<Integer, Set<String>>) loada.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
         assertNull(annotation);

         annotation =
             (Map<Integer, Set<String>>) loadb.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
         assertEquals(1, annotation.size());

         s = new HashSet<String>();
         s.add("path");
         assertEquals(annotation.get(2), s);
     }

     @Test
     public void testPruneWithBag() throws Exception  {
         // filter above foreach
         String query = "a =load 'd.txt' as (id, v:bag{t:(s1,s2,s3)});" +
         "b = filter a by id>10;" +
         "c = foreach b generate id, FLATTEN(v);" +
         "d = foreach c generate id, v::s2;" +
         "store d into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, v:bag{t:(s1,s2,s3)});" +
         "b = filter a by id>10;" +
         "c = foreach b generate id, FLATTEN(v);" +
         "d = foreach c generate id, v::s2;" +
         "store d into 'empty';";
         LogicalPlan expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));
     }

     @Test
     public void testAddForeach() throws Exception  {
         // filter above foreach
         String query = "a =load 'd.txt' as (id, v1, v2);" +
         "b = filter a by v1>10;" +
         "c = foreach b generate id;" +
         "store c into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, v1);" +
         "b = filter a by v1>10;" +
         "c = foreach b generate id;" +
         "store c into 'empty';";
         LogicalPlan expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));

         // join with foreach
         query = "a =load 'd.txt' as (id, v1, v2);" +
         "b = load 'd.txt' as (id, v1, v2);" +
         "c = join a by id, b by id;" +
         "d = filter c by a::v1>b::v1;" +
         "e = foreach d generate a::id;" +
         "store e into 'empty';";
         newLogicalPlan = buildPlan(query);

         optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         query = "a =load 'd.txt' as (id, v1);" +
         "b = load 'd.txt' as (id, v1);" +
         "c = join a by id, b by id;" +
         "d = foreach c generate a::id, a::v1, b::v1;" +
         "e = filter d by a::v1>b::v1;" +
         "f = foreach e generate a::id;" +
         "store f into 'empty';";
         expected = buildPlan(query);

         assertTrue(expected.isEqual(newLogicalPlan));
     }

     @Test
     public void testPruneSubTreeForEach() throws Exception  {
         String query = "a =load 'd.txt' as (id, v1);" +
             "b = group a by id;" +
             "c = foreach b { d = a.v1; " +
                           " e = distinct d; " +
                           " generate group, e; };" +
             "f = foreach c generate group ;" +
             "store f into 'empty';";
         LogicalPlan newLogicalPlan = buildPlan(query);
         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         try {
             optimizer.optimize();
         } catch (Exception e) {
             //PIG-2968 throws ConcurrentModificationException
             e.printStackTrace();
             fail("Unexpected Exception: " + e);
         }
     }

     @Test
     public void testDistinct() throws Exception {
         //Test for bug where distinct wasn't being pruned properly causing union
         //to fail to get a schema since the distinct relation had an incompatible schema
         //with the other relation being unioned.

         String testQuery =
                 "a = load 'd.txt' as (id, v1, v2);" +
                 "b = load 'd.txt' as (id, v1, v2);" +
                 "c = distinct a;" +
                 "d = union b, c;" +
                 "e = foreach d generate id, v1;" +
                 "store e into 'empty';";

         //Generate optimized plan.
         LogicalPlan newLogicalPlan = buildPlan(testQuery);
         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();

         Iterator<Operator> iter = newLogicalPlan.getOperators();
         while (iter.hasNext()) {
             Operator o = iter.next();
             LogicalRelationalOperator lro = (LogicalRelationalOperator)o;
             if (lro == null || lro.getAlias() == null) continue;
             if (lro.getAlias().equals("d")) {
                 assertNotNull(lro.getSchema());
             }
         }
     }

     @Test
     public void testNoAddForeach() throws Exception  {
         // PIG-5055
         // Need to make sure that it does not add foreach
         // that drops all the fields from B2.
         String query = "A = load 'd.txt' as (a0:int, a1:int, a2:int);" +
         "B = load 'd.txt' as (b0:int, b1:int, b2:int);" +
         "B2 = FILTER B by b0 == 0;" +
         "C = join A by (1), B2 by (1) ;" +
         "D = FOREACH C GENERATE A::a1, A::a2;" +
         "store D into 'empty';";

         LogicalPlan newLogicalPlan = buildPlan(query);

         PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
         optimizer.optimize();
         System.err.println(newLogicalPlan);
         Iterator<Operator> iter = newLogicalPlan.getOperators();
         while (iter.hasNext()) {
             Operator o = iter.next();
             LogicalRelationalOperator lro = (LogicalRelationalOperator)o;
             if (lro == null || lro.getAlias() == null) continue;
             if (lro.getAlias().equals("B2")) {
                 assertNotNull(lro.getSchema());
             }
         }
     }

     public class MyPlanOptimizer extends LogicalPlanOptimizer {

         protected MyPlanOptimizer(OperatorPlan p,  int iterations) {
             super(p, iterations, null);
         }

         protected List<Set<Rule>> buildRuleSets() {
             List<Set<Rule>> ls = new ArrayList<Set<Rule>>();

             Rule r = new ColumnMapKeyPrune("ColumnMapKeyPrune");
             Set<Rule> s = new HashSet<Rule>();
             s.add(r);
             ls.add(s);

             r = new AddForEach("AddForEach");
             s = new HashSet<Rule>();
             s.add(r);
             ls.add(s);

             return ls;
         }
     }
 }
	/*
	* 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.assertEquals;
	import static org.junit.Assert.assertNull;
	import static org.junit.Assert.assertNotNull;
	import static org.junit.Assert.assertTrue;
	import static org.junit.Assert.fail;

	import java.util.ArrayList;
	import java.util.HashSet;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.Properties;
	import java.util.Set;

	import org.apache.pig.PigServer;
	import org.apache.pig.impl.PigContext;
	import org.apache.pig.newplan.Operator;
	import org.apache.pig.newplan.OperatorPlan;
	import org.apache.pig.newplan.logical.optimizer.LogicalPlanOptimizer;
	import org.apache.pig.newplan.logical.relational.LOLoad;
	import org.apache.pig.newplan.logical.relational.LogicalPlan;
	import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
	import org.apache.pig.newplan.logical.rules.AddForEach;
	import org.apache.pig.newplan.logical.rules.ColumnMapKeyPrune;
	import org.apache.pig.newplan.logical.rules.MapKeysPruneHelper;
	import org.apache.pig.newplan.optimizer.PlanOptimizer;
	import org.apache.pig.newplan.optimizer.Rule;
	import org.junit.Before;
	import org.junit.Test;

	public class TestNewPlanColumnPrune {
	LogicalPlan plan = null;
	PigContext pc;

	@Before
	public void setUp() throws Exception {
	pc = new PigContext(Util.getLocalTestMode(), new Properties());
	}
	private LogicalPlan buildPlan(String query) throws Exception{
	PigServer pigServer = new PigServer( pc );
	return Util.buildLp(pigServer, query);
	}

	@Test
	public void testNoPrune() throws Exception {
	// no foreach
	String query = "a = load 'd.txt' as (id, v1, v2);" +
	"b = filter a by v1==NULL;" +
	"store b into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt' as (id, v1, v2);" +
	"b = filter a by v1==NULL;" +
	"store b into 'empty';";
	LogicalPlan expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));

	// no schema
	query = "a = load 'd.txt';" +
	"b = foreach a generate $0, $1;" +
	"store b into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt';"+
	"b = foreach a generate $0, $1;"+
	"store b into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));
	}

	@Test
	public void testPrune() throws Exception {
	// only foreach
	String query = "a = load 'd.txt' as (id, v1, v2);" +
	"b = foreach a generate id;"+
	"store b into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt' as (id);" +
	"b = foreach a generate id;"+
	"store b into 'empty';";
	LogicalPlan expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));

	// with filter
	query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);"+
	"b = filter a by v1 != NULL AND (v2+v3)<100;"+
	"c = foreach b generate id;"+
	"store c into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt' as (id, v1, v3, v2);" +
	"b = filter a by v1 != NULL AND (v2+v3)<100;" +
	"c = foreach b generate id;" +
	"store c into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with 2 foreach
	query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v2, v5, v4;" +
	"c = foreach b generate v5, v4;" +
	"store c into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt' as (v5, v4);" +
	"b = foreach a generate v5, v4;" +
	"c = foreach b generate v5, v4;" +
	"store c into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with 2 foreach
	query = "a = load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate id, v1, v5, v3, v4;" +
	"c = foreach b generate v5, v4;" +
	"store c into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a = load 'd.txt' as (v5, v4);" +
	"b = foreach a generate v5, v4;" +
	"c = foreach b generate v5, v4;" +
	"store c into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with 2 foreach and filter in between
	query = "a =load 'd.txt' as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v2, v5, v4;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5, v4;" +
	"store d into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (v5, v4, v2);" +
	"b = foreach a generate v2, v5, v4;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5, v4;" +
	"store d into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with 2 foreach after join
	query = "a =load 'd.txt' as (id, v1, v2, v3);" +
	"b = load 'c.txt' as (id, v4, v5, v6);" +
	"c = join a by id, b by id;" +
	"d = foreach c generate a::id, v5, v3, v4;" +
	"store d into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, v3);" +
	"b = load 'c.txt' as (id, v4, v5);" +
	"c = join a by id, b by id;" +
	"d = foreach c generate a::id, v5, v3, v4;" +
	"store d into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with BinStorage, insert foreach after load
	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"c = filter a by v2 != NULL;" +
	"d = foreach c generate v5, v4;" +
	"store d into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v5, v4, v2;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5, v4;" +
	"store d into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with BinStorage, not to insert foreach after load if there is already one
	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v5, v4, v2;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5;" +
	"store d into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v5, v2;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5;" +
	"store d into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));

	// with BinStorage, not to insert foreach after load if there is already one
	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v5, v4, v2, 10;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5;" +
	"store d into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' using BinStorage() as (id, v1, v5, v3, v4, v2);" +
	"b = foreach a generate v5, v2, 10;" +
	"c = filter b by v2 != NULL;" +
	"d = foreach c generate v5;" +
	"store d into 'empty';";
	expected = buildPlan(query);
	assertTrue(expected.isEqual(newLogicalPlan));
	}

	@Test
	@SuppressWarnings("unchecked")
	public void testPruneWithMapKey() throws Exception {
	// only foreach
	String query = "a =load 'd.txt' as (id, v1, m:map[]);" +
	"b = foreach a generate id, m#'path';" +
	"store b into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, m:map[]);" +
	"b = foreach a generate id, m#'path';" +
	"store b into 'empty';";
	LogicalPlan expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));

	LOLoad op = (LOLoad)newLogicalPlan.getSources().get(0);
	Map<Integer,Set<String>> annotation =
	(Map<Integer, Set<String>>) op.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
	assertEquals(1, annotation.size());
	Set<String> s = new HashSet<String>();
	s.add("path");
	assertEquals(annotation.get(2), s);

	// foreach with join
	query = "a =load 'd.txt' as (id, v1, m:map[]);" +
	"b = load 'd.txt' as (id, v1, m:map[]);" +
	"c = join a by id, b by id;" +
	"d = filter c by a::m#'path' != NULL;" +
	"e = foreach d generate a::id, b::id, b::m#'path', a::m;" +
	"store e into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, m:map[]);" +
	"b = load 'd.txt' as (id, m:map[]);" +
	"c = join a by id, b by id;" +
	"d = filter c by a::m#'path' != NULL;" +
	"e = foreach d generate a::id, b::id, b::m#'path', a::m;" +
	"store e into 'empty';";
	expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));

	List<Operator> ll = newLogicalPlan.getSources();
	assertEquals(2, ll.size());
	LOLoad loada = null;
	LOLoad loadb = null;
	for(Operator opp: ll) {
	if (((LogicalRelationalOperator)opp).getAlias().equals("a")) {
	loada = (LOLoad)opp;
	continue;
	}

	if (((LogicalRelationalOperator)opp).getAlias().equals("b")) {
	loadb = (LOLoad)opp;
	continue;
	}
	}

	annotation =
	(Map<Integer, Set<String>>) loada.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
	assertNull(annotation);

	annotation =
	(Map<Integer, Set<String>>) loadb.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
	assertEquals(1, annotation.size());

	s = new HashSet<String>();
	s.add("path");
	assertEquals(annotation.get(2), s);
	}

	@Test
	public void testPruneWithBag() throws Exception {
	// filter above foreach
	String query = "a =load 'd.txt' as (id, v:bag{t:(s1,s2,s3)});" +
	"b = filter a by id>10;" +
	"c = foreach b generate id, FLATTEN(v);" +
	"d = foreach c generate id, v::s2;" +
	"store d into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, v:bag{t:(s1,s2,s3)});" +
	"b = filter a by id>10;" +
	"c = foreach b generate id, FLATTEN(v);" +
	"d = foreach c generate id, v::s2;" +
	"store d into 'empty';";
	LogicalPlan expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));
	}

	@Test
	public void testAddForeach() throws Exception {
	// filter above foreach
	String query = "a =load 'd.txt' as (id, v1, v2);" +
	"b = filter a by v1>10;" +
	"c = foreach b generate id;" +
	"store c into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, v1);" +
	"b = filter a by v1>10;" +
	"c = foreach b generate id;" +
	"store c into 'empty';";
	LogicalPlan expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));

	// join with foreach
	query = "a =load 'd.txt' as (id, v1, v2);" +
	"b = load 'd.txt' as (id, v1, v2);" +
	"c = join a by id, b by id;" +
	"d = filter c by a::v1>b::v1;" +
	"e = foreach d generate a::id;" +
	"store e into 'empty';";
	newLogicalPlan = buildPlan(query);

	optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	query = "a =load 'd.txt' as (id, v1);" +
	"b = load 'd.txt' as (id, v1);" +
	"c = join a by id, b by id;" +
	"d = foreach c generate a::id, a::v1, b::v1;" +
	"e = filter d by a::v1>b::v1;" +
	"f = foreach e generate a::id;" +
	"store f into 'empty';";
	expected = buildPlan(query);

	assertTrue(expected.isEqual(newLogicalPlan));
	}

	@Test
	public void testPruneSubTreeForEach() throws Exception {
	String query = "a =load 'd.txt' as (id, v1);" +
	"b = group a by id;" +
	"c = foreach b { d = a.v1; " +
	" e = distinct d; " +
	" generate group, e; };" +
	"f = foreach c generate group ;" +
	"store f into 'empty';";
	LogicalPlan newLogicalPlan = buildPlan(query);
	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	try {
	optimizer.optimize();
	} catch (Exception e) {
	//PIG-2968 throws ConcurrentModificationException
	e.printStackTrace();
	fail("Unexpected Exception: " + e);
	}
	}

	@Test
	public void testDistinct() throws Exception {
	//Test for bug where distinct wasn't being pruned properly causing union
	//to fail to get a schema since the distinct relation had an incompatible schema
	//with the other relation being unioned.

	String testQuery =
	"a = load 'd.txt' as (id, v1, v2);" +
	"b = load 'd.txt' as (id, v1, v2);" +
	"c = distinct a;" +
	"d = union b, c;" +
	"e = foreach d generate id, v1;" +
	"store e into 'empty';";

	//Generate optimized plan.
	LogicalPlan newLogicalPlan = buildPlan(testQuery);
	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();

	Iterator<Operator> iter = newLogicalPlan.getOperators();
	while (iter.hasNext()) {
	Operator o = iter.next();
	LogicalRelationalOperator lro = (LogicalRelationalOperator)o;
	if (lro == null \|\| lro.getAlias() == null) continue;
	if (lro.getAlias().equals("d")) {
	assertNotNull(lro.getSchema());
	}
	}
	}

	@Test
	public void testNoAddForeach() throws Exception {
	// PIG-5055
	// Need to make sure that it does not add foreach
	// that drops all the fields from B2.
	String query = "A = load 'd.txt' as (a0:int, a1:int, a2:int);" +
	"B = load 'd.txt' as (b0:int, b1:int, b2:int);" +
	"B2 = FILTER B by b0 == 0;" +
	"C = join A by (1), B2 by (1) ;" +
	"D = FOREACH C GENERATE A::a1, A::a2;" +
	"store D into 'empty';";

	LogicalPlan newLogicalPlan = buildPlan(query);

	PlanOptimizer optimizer = new MyPlanOptimizer(newLogicalPlan, 3);
	optimizer.optimize();
	System.err.println(newLogicalPlan);
	Iterator<Operator> iter = newLogicalPlan.getOperators();
	while (iter.hasNext()) {
	Operator o = iter.next();
	LogicalRelationalOperator lro = (LogicalRelationalOperator)o;
	if (lro == null \|\| lro.getAlias() == null) continue;
	if (lro.getAlias().equals("B2")) {
	assertNotNull(lro.getSchema());
	}
	}
	}

	public class MyPlanOptimizer extends LogicalPlanOptimizer {

	protected MyPlanOptimizer(OperatorPlan p, int iterations) {
	super(p, iterations, null);
	}

	protected List<Set<Rule>> buildRuleSets() {
	List<Set<Rule>> ls = new ArrayList<Set<Rule>>();

	Rule r = new ColumnMapKeyPrune("ColumnMapKeyPrune");
	Set<Rule> s = new HashSet<Rule>();
	s.add(r);
	ls.add(s);

	r = new AddForEach("AddForEach");
	s = new HashSet<Rule>();
	s.add(r);
	ls.add(s);

	return ls;
	}
	}
	}