test/org/apache/pig/test/TestCombiner.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.assertFalse;
 import static org.junit.Assert.assertTrue;

 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Properties;

 import junit.framework.Assert;

 import org.apache.pig.EvalFunc;
 import org.apache.pig.ExecType;
 import org.apache.pig.PigServer;
 import org.apache.pig.ResourceSchema.ResourceFieldSchema;
 import org.apache.pig.builtin.PigStorage;
 import org.apache.pig.data.DataBag;
 import org.apache.pig.data.DataType;
 import org.apache.pig.data.DefaultDataBag;
 import org.apache.pig.data.Tuple;
 import org.apache.pig.impl.PigContext;
 import org.apache.pig.impl.io.FileLocalizer;
 import org.apache.pig.impl.util.Utils;
 import org.apache.pig.newplan.logical.relational.LogicalSchema;
 import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;
 import org.junit.AfterClass;
 import org.junit.Test;

 public class TestCombiner {

     static MiniCluster cluster = MiniCluster.buildCluster();

     @AfterClass
     public static void oneTimeTearDown() throws Exception {
         cluster.shutDown();
     }


     @Test
     public void testSuccessiveUserFuncs1() throws Exception{
         String query = "a = load 'students.txt' as (c1,c2,c3,c4); " +
                        "c = group a by c2; " +
                        "f = foreach c generate COUNT(org.apache.pig.builtin.Distinct($1.$2)); " +
                        "store f into 'out';";
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         PigContext pc = pigServer.getPigContext();
         assertTrue((Util.buildMRPlan(Util.buildPp(pigServer,query),pc).getRoots().get(0).combinePlan.isEmpty()));
     }

     @Test
     public void testSuccessiveUserFuncs2() throws Exception {
         String dummyUDF = JiraPig1030.class.getName();
     	String query = "a = load 'students.txt' as (c1,c2,c3,c4); " +
                        "c = group a by c2; " +
                        "f = foreach c generate COUNT(" + dummyUDF + "" +
                		   "(org.apache.pig.builtin.Distinct($1.$2),"+dummyUDF+"())); " +
                		   "store f into 'out';";
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         PigContext pc = pigServer.getPigContext();
         assertTrue((Util.buildMRPlan(Util.buildPp(pigServer,query),pc).getRoots().get(0).combinePlan.isEmpty()));
     }

     @Test
     public void testOnCluster() throws Exception {
         // run the test on cluster
         String inputFileName = runTest(new PigServer(
                 ExecType.MAPREDUCE, cluster.getProperties()));
         Util.deleteFile(cluster, inputFileName);

     }

     /* (non-Javadoc)
      * @see junit.framework.TestCase#setUp()
      */
     protected void setUp() throws Exception {
         // cause a re initialization of FileLocalizer's
         // internal state before each test run
         // A previous test might have been in a different
         // mode than the test which is about to run. To
         // ensure each test runs correctly in it's exectype
         // mode, let's re initialize.
         FileLocalizer.setInitialized(false);
     }

     @Test
     public void testLocal() throws Exception {
         // run the test locally
         FileLocalizer.deleteTempFiles();
         runTest(new PigServer(ExecType.LOCAL, new Properties()));
         FileLocalizer.deleteTempFiles();
     }


     private String  runTest(PigServer pig) throws IOException {
         List<String> inputLines = new ArrayList<String>();
         inputLines.add("a,b,1");
         inputLines.add("a,b,1");
         inputLines.add("a,c,1");
         String inputFileName = loadWithTestLoadFunc("A", pig, inputLines);

         pig.registerQuery("B = group A by ($0, $1);");
         pig.registerQuery("C = foreach B generate flatten(group), COUNT($1);");
         Iterator<Tuple> resultIterator = pig.openIterator("C");
         Tuple tuple = resultIterator.next();
         assertEquals("(a,b,2)", tuple.toString());
         tuple = resultIterator.next();
         assertEquals("(a,c,1)", tuple.toString());

         return inputFileName;
     }

     private String loadWithTestLoadFunc(String loadAlias, PigServer pig,
             List<String> inputLines) throws IOException {
         File inputFile = File.createTempFile("test", "txt");
         inputFile.deleteOnExit();
         String inputFileName = inputFile.getAbsolutePath();
         if(pig.getPigContext().getExecType() == ExecType.LOCAL) {
             PrintStream ps = new PrintStream(new FileOutputStream(inputFile));
             for (String line : inputLines) {
                 ps.println(line);
             }
             ps.close();
         } else {
             inputFileName = Util.removeColon(inputFileName);
             Util.createInputFile(cluster, inputFileName, inputLines.toArray(new String[] {}));
         }
         pig.registerQuery(loadAlias + " = load '"
                 + Util.encodeEscape(inputFileName) + "' using "
                 + PigStorage.class.getName() + "(',');");
         return inputFileName;
     }

     @Test
     public void testNoCombinerUse() {
         // To simulate this, we will have two input files
         // with exactly one input record - this should result
         // in two map tasks and each would process only one record
         // hence the combiner would not get called.
     }

     @Test
     public void testMultiCombinerUse() throws Exception {
         // test the scenario where the combiner is called multiple
         // times - this can happen when the output of the map > io.sort.mb
         // let's set the io.sort.mb to 1MB and > 1 MB map data.
         String[] input = new String[500*1024];
         for(int i = 0; i < input.length; i++) {
             if(i % 2 == 0) {
                 input[i] = Integer.toString(1);
             } else {
                 input[i] = Integer.toString(0);
             }
         }
         Util.createInputFile(cluster, "MultiCombinerUseInput.txt", input);
         Properties props = cluster.getProperties();
         props.setProperty("io.sort.mb", "1");
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, props);
         pigServer.registerQuery("a = load 'MultiCombinerUseInput.txt' as (x:int);");
         pigServer.registerQuery("b = group a all;");
         pigServer.registerQuery("c = foreach b generate COUNT(a), SUM(a.$0), " +
         		"MIN(a.$0), MAX(a.$0), AVG(a.$0), ((double)SUM(a.$0))/COUNT(a.$0)," +
         		" COUNT(a.$0) + SUM(a.$0) +  MAX(a.$0);");

         // make sure there is a combine plan in the explain output
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         PrintStream ps = new PrintStream(baos);
         pigServer.explain("c", ps);
         assertTrue(baos.toString().matches("(?si).*combine plan.*"));

         Iterator<Tuple> it = pigServer.openIterator("c");
         Tuple t = it.next();
         assertEquals(512000L, t.get(0));
         assertEquals(256000L, t.get(1));
         assertEquals(0, t.get(2));
         assertEquals(1, t.get(3));
         assertEquals(0.5, t.get(4));
         assertEquals(0.5, t.get(5));
         assertEquals(512000L + 256000L + 1, t.get(6));

         assertFalse(it.hasNext());
         Util.deleteFile(cluster, "MultiCombinerUseInput.txt");
     }

     @Test
     public void testDistinctAggs1() throws Exception {
         // test the use of combiner for distinct aggs:
         String input[] = {
                 "pig1\t18\t2.1",
                 "pig2\t24\t3.3",
                 "pig5\t45\t2.4",
                 "pig1\t18\t2.1",
                 "pig1\t19\t2.1",
                 "pig2\t24\t4.5",
                 "pig1\t20\t3.1" };

         Util.createInputFile(cluster, "distinctAggs1Input.txt", input);
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         pigServer.registerQuery("a = load 'distinctAggs1Input.txt' as (name:chararray, age:int, gpa:double);");
         pigServer.registerQuery("b = group a by name;");
         pigServer.registerQuery("c = foreach b  {" +
         		"        x = distinct a.age;" +
         		"        y = distinct a.gpa;" +
         		"        z = distinct a;" +
         		"        generate group, COUNT(x), SUM(x.age), SUM(y.gpa), SUM(a.age), " +
         		"                       SUM(a.gpa), COUNT(z.age), COUNT(z), SUM(z.age);};");

         // make sure there is a combine plan in the explain output
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         PrintStream ps = new PrintStream(baos);
         pigServer.explain("c", ps);
         assertTrue(baos.toString().matches("(?si).*combine plan.*"));

         HashMap<String, Object[]> results = new HashMap<String, Object[]>();
         results.put("pig1", new Object[] {"pig1",3L,57L,5.2,75L,9.4,3L,3L,57L});
         results.put("pig2", new Object[] {"pig2",1L,24L,7.8,48L,7.8,2L,2L,48L});
         results.put("pig5", new Object[] {"pig5",1L,45L,2.4,45L,2.4,1L,1L,45L});
         Iterator<Tuple> it = pigServer.openIterator("c");
         while(it.hasNext()) {
             Tuple t = it.next();
             List<Object> fields = t.getAll();
             Object[] expected = results.get((String)fields.get(0));
             int i = 0;
             for (Object field : fields) {
                 assertEquals(expected[i++], field);
             }
         }
         Util.deleteFile(cluster, "distinctAggs1Input.txt");

     }

     @Test
     public void testGroupElements() throws Exception {
         // test use of combiner when group elements are accessed in the foreach
         String input[] = {
                 "ABC\t1\ta\t1",
                 "ABC\t1\tb\t2",
                 "ABC\t1\ta\t3",
                 "ABC\t2\tb\t4",
                 "DEF\t1\td\t1",
                 "XYZ\t1\tx\t2"
         };

         Util.createInputFile(cluster, "testGroupElements.txt", input);
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         pigServer.registerQuery("a = load 'testGroupElements.txt' as (str:chararray, num1:int, alph : chararray, num2 : int);");
         pigServer.registerQuery("b = group a by (str, num1);");

         //check if combiner is present or not for various forms of foreach
         pigServer.registerQuery("c = foreach b  generate flatten(group), COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", true);

         pigServer.registerQuery("c = foreach b  generate group, COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", true);

         // projecting bag - combiner should not be used
         pigServer.registerQuery("c = foreach b  generate group, a,  COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", false);

         // projecting bag - combiner should not be used
         pigServer.registerQuery("c = foreach b  generate group, a.num2,  COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", false);

         pigServer.registerQuery("c = foreach b  generate group.$0, group.$1, COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", true);

         pigServer.registerQuery("c = foreach b  generate group.$0, group.$1 + COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", true);

         pigServer.registerQuery("c = foreach b  generate group.str, group.$1, COUNT(a.alph), SUM(a.num2); ");
         checkCombinerUsed(pigServer, "c", true);

         pigServer.registerQuery("c = foreach b  generate group.str, group.$1, COUNT(a.alph), SUM(a.num2), " +
         		" (group.num1 == 1 ? (COUNT(a.num2) + 1)  : (SUM(a.num2) + 10)) ; ");
         checkCombinerUsed(pigServer, "c", true);

         List<Tuple> expectedRes =
             Util.getTuplesFromConstantTupleStrings(
                     new String[] {
                             "('ABC',1,3L,6L,4L)",
                             "('ABC',2,1L,4L,14L)",
                             "('DEF',1,1L,1L,2L)",
                             "('XYZ',1,1L,2L,2L)",
                     });

         Iterator<Tuple> it = pigServer.openIterator("c");
         Util.checkQueryOutputsAfterSort(it, expectedRes);

         Util.deleteFile(cluster, "distinctAggs1Input.txt");

     }

     @Test
     public void testGroupByLimit() throws Exception {
         // test use of combiner when group elements are accessed in the foreach
         String input[] = {
                 "ABC 1",
                 "ABC 2",
                 "DEF 1",
                 "XYZ 1",
                 "XYZ 2",
                 "XYZ 3",
         };

         Util.createInputFile(cluster, "testGroupLimit.txt", input);
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         pigServer.registerQuery("a = load 'testGroupLimit.txt'  using PigStorage(' ') " +
         		"as (str:chararray, num1:int) ;");
         pigServer.registerQuery("b = group a by str;");


         pigServer.registerQuery("c = foreach b  generate group, COUNT(a.num1) ; ");

         //check if combiner is present
         pigServer.registerQuery("d = limit c 2 ; ");
         checkCombinerUsed(pigServer, "d", true);

         List<Tuple> expectedRes =
             Util.getTuplesFromConstantTupleStrings(
                     new String[] {
                             "('ABC',2L)",
                             "('DEF',1L)",
                     });

         Iterator<Tuple> it = pigServer.openIterator("d");
         Util.checkQueryOutputsAfterSort(it, expectedRes);


     }

     private void checkCombinerUsed(PigServer pigServer, String string, boolean combineExpected)
     throws IOException {
         // make sure there is a combine plan in the explain output
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         PrintStream ps = new PrintStream(baos);
         pigServer.explain("c", ps);
         boolean combinerFound = baos.toString().matches("(?si).*combine plan.*");
         System.out.println(baos.toString());
         assertEquals("is combiner present as expected", combineExpected, combinerFound);
     }


     @Test
     public void testDistinctNoCombiner() throws Exception {
         // test that combiner is NOT invoked when
         // one of the elements in the foreach generate
         // is a distinct() as the leaf
         String input[] = {
                 "pig1\t18\t2.1",
                 "pig2\t24\t3.3",
                 "pig5\t45\t2.4",
                 "pig1\t18\t2.1",
                 "pig1\t19\t2.1",
                 "pig2\t24\t4.5",
                 "pig1\t20\t3.1" };

         Util.createInputFile(cluster, "distinctNoCombinerInput.txt", input);
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         pigServer.registerQuery("a = load 'distinctNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
         pigServer.registerQuery("b = group a by name;");
         pigServer.registerQuery("c = foreach b  {" +
                 "        z = distinct a;" +
                 "        generate group, z, SUM(a.age), SUM(a.gpa);};");

         // make sure there is a combine plan in the explain output
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         PrintStream ps = new PrintStream(baos);
         pigServer.explain("c", ps);
         assertFalse(baos.toString().matches("(?si).*combine plan.*"));

         HashMap<String, Object[]> results = new HashMap<String, Object[]>();
         results.put("pig1", new Object[] {"pig1","bag-place-holder",75L,9.4});
         results.put("pig2", new Object[] {"pig2","bag-place-holder",48L,7.8});
         results.put("pig5", new Object[] {"pig5","bag-place-holder",45L,2.4});
         Iterator<Tuple> it = pigServer.openIterator("c");
         while(it.hasNext()) {
             Tuple t = it.next();
             List<Object> fields = t.getAll();
             Object[] expected = results.get((String)fields.get(0));
             int i = 0;
             for (Object field : fields) {
                 if(i == 1) {
                     // ignore the second field which is a bag
                     // for comparison here
                     continue;
                 }
                 assertEquals(expected[i++], field);
             }
         }
         Util.deleteFile(cluster, "distinctNoCombinerInput.txt");

     }

     @Test
     public void testForEachNoCombiner() throws Exception {
         // test that combiner is NOT invoked when
         // one of the elements in the foreach generate
         // has a foreach in the plan without a distinct agg
         String input[] = {
                 "pig1\t18\t2.1",
                 "pig2\t24\t3.3",
                 "pig5\t45\t2.4",
                 "pig1\t18\t2.1",
                 "pig1\t19\t2.1",
                 "pig2\t24\t4.5",
                 "pig1\t20\t3.1" };

         Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);
         PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
         pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
         pigServer.registerQuery("b = group a by name;");
         pigServer.registerQuery("c = foreach b  {" +
                 "        z = a.age;" +
                 "        generate group, z, SUM(a.age), SUM(a.gpa);};");

         // make sure there is a combine plan in the explain output
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         PrintStream ps = new PrintStream(baos);
         pigServer.explain("c", ps);
         assertFalse(baos.toString().matches("(?si).*combine plan.*"));

         HashMap<String, Object[]> results = new HashMap<String, Object[]>();
         results.put("pig1", new Object[] {"pig1","bag-place-holder",75L,9.4});
         results.put("pig2", new Object[] {"pig2","bag-place-holder",48L,7.8});
         results.put("pig5", new Object[] {"pig5","bag-place-holder",45L,2.4});
         Iterator<Tuple> it = pigServer.openIterator("c");
         while(it.hasNext()) {
             Tuple t = it.next();
             List<Object> fields = t.getAll();
             Object[] expected = results.get((String)fields.get(0));
             int i = 0;
             for (Object field : fields) {
                 if(i == 1) {
                     // ignore the second field which is a bag
                     // for comparison here
                     continue;
                 }
                 assertEquals(expected[i++], field);
             }
         }
         Util.deleteFile(cluster, "forEachNoCombinerInput.txt");

     }

     @Test
     public void testJiraPig746() {
         // test that combiner is NOT invoked when
         // one of the elements in the foreach generate
         // has a foreach in the plan without a distinct agg
         String input[] = {
                 "pig1\t18\t2.1",
                 "pig2\t24\t3.3",
                 "pig5\t45\t2.4",
                 "pig1\t18\t2.1",
                 "pig1\t19\t2.1",
                 "pig2\t24\t4.5",
                 "pig1\t20\t3.1" };

         String expected[] = {
                 "(pig1,75,{(pig1,18,2.1),(pig1,18,2.1),(pig1,19,2.1),(pig1,20,3.1)})",
                 "(pig2,48,{(pig2,24,3.3),(pig2,24,4.5)})",
                 "(pig5,45,{(pig5,45,2.4)})"
         };

         try {
             Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);

             PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
             pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
             pigServer.registerQuery("b = group a by name;");
             pigServer.registerQuery("c = foreach b generate group, SUM(a.age), a;");

             // make sure there isn't a combine plan in the explain output
             ByteArrayOutputStream baos = new ByteArrayOutputStream();
             PrintStream ps = new PrintStream(baos);
             pigServer.explain("c", ps);
             assertFalse(baos.toString().matches("(?si).*combine plan.*"));

             Iterator<Tuple> it = pigServer.openIterator("c");
             Util.checkQueryOutputsAfterSortRecursive(it, expected, "group:chararray,age:long,b:{t:(name:chararray,age:int,gpa:double)}");
         } catch (IOException e) {
             e.printStackTrace();
             Assert.fail();
         } finally {
             try {
                 Util.deleteFile(cluster, "forEachNoCombinerInput.txt");
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }

     public static class JiraPig1030 extends EvalFunc<DataBag> {

         public DataBag exec(Tuple input) throws IOException {
             return new DefaultDataBag();
         }
     }

     @Test
     public void testJiraPig1030() {
         // test that combiner is NOT invoked when
         // one of the elements in the foreach generate
         // has a non-algebraic UDF that have multiple inputs
         // (one of them is distinct).

         String input[] = {
                 "pig1\t18\t2.1",
                 "pig2\t24\t3.3",
                 "pig5\t45\t2.4",
                 "pig1\t18\t2.1",
                 "pig1\t19\t2.1",
                 "pig2\t24\t4.5",
                 "pig1\t20\t3.1" };

         try {
             Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);
             PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
             pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
             pigServer.registerQuery("b = group a all;");
             pigServer.registerQuery("c = foreach b  {" +
                     "        d = distinct a.age;" +
                     "        generate group, " + JiraPig1030.class.getName() + "(d, 0);};");

             // make sure there isn't a combine plan in the explain output
             ByteArrayOutputStream baos = new ByteArrayOutputStream();
             PrintStream ps = new PrintStream(baos);
             pigServer.explain("c", ps);
             assertFalse(baos.toString().matches("(?si).*combine plan.*"));
         } catch (Exception e) {
             e.printStackTrace();
             Assert.fail();
         } finally {
             try {
                 Util.deleteFile(cluster, "forEachNoCombinerInput.txt");
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }

 }
	/*
	* 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.assertFalse;
	import static org.junit.Assert.assertTrue;

	import java.io.ByteArrayOutputStream;
	import java.io.File;
	import java.io.FileOutputStream;
	import java.io.IOException;
	import java.io.PrintStream;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Properties;

	import junit.framework.Assert;

	import org.apache.pig.EvalFunc;
	import org.apache.pig.ExecType;
	import org.apache.pig.PigServer;
	import org.apache.pig.ResourceSchema.ResourceFieldSchema;
	import org.apache.pig.builtin.PigStorage;
	import org.apache.pig.data.DataBag;
	import org.apache.pig.data.DataType;
	import org.apache.pig.data.DefaultDataBag;
	import org.apache.pig.data.Tuple;
	import org.apache.pig.impl.PigContext;
	import org.apache.pig.impl.io.FileLocalizer;
	import org.apache.pig.impl.util.Utils;
	import org.apache.pig.newplan.logical.relational.LogicalSchema;
	import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;
	import org.junit.AfterClass;
	import org.junit.Test;

	public class TestCombiner {

	static MiniCluster cluster = MiniCluster.buildCluster();

	@AfterClass
	public static void oneTimeTearDown() throws Exception {
	cluster.shutDown();
	}


	@Test
	public void testSuccessiveUserFuncs1() throws Exception{
	String query = "a = load 'students.txt' as (c1,c2,c3,c4); " +
	"c = group a by c2; " +
	"f = foreach c generate COUNT(org.apache.pig.builtin.Distinct($1.$2)); " +
	"store f into 'out';";
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	PigContext pc = pigServer.getPigContext();
	assertTrue((Util.buildMRPlan(Util.buildPp(pigServer,query),pc).getRoots().get(0).combinePlan.isEmpty()));
	}

	@Test
	public void testSuccessiveUserFuncs2() throws Exception {
	String dummyUDF = JiraPig1030.class.getName();
	String query = "a = load 'students.txt' as (c1,c2,c3,c4); " +
	"c = group a by c2; " +
	"f = foreach c generate COUNT(" + dummyUDF + "" +
	"(org.apache.pig.builtin.Distinct($1.$2),"+dummyUDF+"())); " +
	"store f into 'out';";
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	PigContext pc = pigServer.getPigContext();
	assertTrue((Util.buildMRPlan(Util.buildPp(pigServer,query),pc).getRoots().get(0).combinePlan.isEmpty()));
	}

	@Test
	public void testOnCluster() throws Exception {
	// run the test on cluster
	String inputFileName = runTest(new PigServer(
	ExecType.MAPREDUCE, cluster.getProperties()));
	Util.deleteFile(cluster, inputFileName);

	}

	/* (non-Javadoc)
	* @see junit.framework.TestCase#setUp()
	*/
	protected void setUp() throws Exception {
	// cause a re initialization of FileLocalizer's
	// internal state before each test run
	// A previous test might have been in a different
	// mode than the test which is about to run. To
	// ensure each test runs correctly in it's exectype
	// mode, let's re initialize.
	FileLocalizer.setInitialized(false);
	}

	@Test
	public void testLocal() throws Exception {
	// run the test locally
	FileLocalizer.deleteTempFiles();
	runTest(new PigServer(ExecType.LOCAL, new Properties()));
	FileLocalizer.deleteTempFiles();
	}


	private String runTest(PigServer pig) throws IOException {
	List<String> inputLines = new ArrayList<String>();
	inputLines.add("a,b,1");
	inputLines.add("a,b,1");
	inputLines.add("a,c,1");
	String inputFileName = loadWithTestLoadFunc("A", pig, inputLines);

	pig.registerQuery("B = group A by ($0, $1);");
	pig.registerQuery("C = foreach B generate flatten(group), COUNT($1);");
	Iterator<Tuple> resultIterator = pig.openIterator("C");
	Tuple tuple = resultIterator.next();
	assertEquals("(a,b,2)", tuple.toString());
	tuple = resultIterator.next();
	assertEquals("(a,c,1)", tuple.toString());

	return inputFileName;
	}

	private String loadWithTestLoadFunc(String loadAlias, PigServer pig,
	List<String> inputLines) throws IOException {
	File inputFile = File.createTempFile("test", "txt");
	inputFile.deleteOnExit();
	String inputFileName = inputFile.getAbsolutePath();
	if(pig.getPigContext().getExecType() == ExecType.LOCAL) {
	PrintStream ps = new PrintStream(new FileOutputStream(inputFile));
	for (String line : inputLines) {
	ps.println(line);
	}
	ps.close();
	} else {
	inputFileName = Util.removeColon(inputFileName);
	Util.createInputFile(cluster, inputFileName, inputLines.toArray(new String[] {}));
	}
	pig.registerQuery(loadAlias + " = load '"
	+ Util.encodeEscape(inputFileName) + "' using "
	+ PigStorage.class.getName() + "(',');");
	return inputFileName;
	}

	@Test
	public void testNoCombinerUse() {
	// To simulate this, we will have two input files
	// with exactly one input record - this should result
	// in two map tasks and each would process only one record
	// hence the combiner would not get called.
	}

	@Test
	public void testMultiCombinerUse() throws Exception {
	// test the scenario where the combiner is called multiple
	// times - this can happen when the output of the map > io.sort.mb
	// let's set the io.sort.mb to 1MB and > 1 MB map data.
	String[] input = new String[500*1024];
	for(int i = 0; i < input.length; i++) {
	if(i % 2 == 0) {
	input[i] = Integer.toString(1);
	} else {
	input[i] = Integer.toString(0);
	}
	}
	Util.createInputFile(cluster, "MultiCombinerUseInput.txt", input);
	Properties props = cluster.getProperties();
	props.setProperty("io.sort.mb", "1");
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, props);
	pigServer.registerQuery("a = load 'MultiCombinerUseInput.txt' as (x:int);");
	pigServer.registerQuery("b = group a all;");
	pigServer.registerQuery("c = foreach b generate COUNT(a), SUM(a.$0), " +
	"MIN(a.$0), MAX(a.$0), AVG(a.$0), ((double)SUM(a.$0))/COUNT(a.$0)," +
	" COUNT(a.$0) + SUM(a.$0) + MAX(a.$0);");

	// make sure there is a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertTrue(baos.toString().matches("(?si).combine plan."));

	Iterator<Tuple> it = pigServer.openIterator("c");
	Tuple t = it.next();
	assertEquals(512000L, t.get(0));
	assertEquals(256000L, t.get(1));
	assertEquals(0, t.get(2));
	assertEquals(1, t.get(3));
	assertEquals(0.5, t.get(4));
	assertEquals(0.5, t.get(5));
	assertEquals(512000L + 256000L + 1, t.get(6));

	assertFalse(it.hasNext());
	Util.deleteFile(cluster, "MultiCombinerUseInput.txt");
	}

	@Test
	public void testDistinctAggs1() throws Exception {
	// test the use of combiner for distinct aggs:
	String input[] = {
	"pig1\t18\t2.1",
	"pig2\t24\t3.3",
	"pig5\t45\t2.4",
	"pig1\t18\t2.1",
	"pig1\t19\t2.1",
	"pig2\t24\t4.5",
	"pig1\t20\t3.1" };

	Util.createInputFile(cluster, "distinctAggs1Input.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'distinctAggs1Input.txt' as (name:chararray, age:int, gpa:double);");
	pigServer.registerQuery("b = group a by name;");
	pigServer.registerQuery("c = foreach b {" +
	" x = distinct a.age;" +
	" y = distinct a.gpa;" +
	" z = distinct a;" +
	" generate group, COUNT(x), SUM(x.age), SUM(y.gpa), SUM(a.age), " +
	" SUM(a.gpa), COUNT(z.age), COUNT(z), SUM(z.age);};");

	// make sure there is a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertTrue(baos.toString().matches("(?si).combine plan."));

	HashMap<String, Object[]> results = new HashMap<String, Object[]>();
	results.put("pig1", new Object[] {"pig1",3L,57L,5.2,75L,9.4,3L,3L,57L});
	results.put("pig2", new Object[] {"pig2",1L,24L,7.8,48L,7.8,2L,2L,48L});
	results.put("pig5", new Object[] {"pig5",1L,45L,2.4,45L,2.4,1L,1L,45L});
	Iterator<Tuple> it = pigServer.openIterator("c");
	while(it.hasNext()) {
	Tuple t = it.next();
	List<Object> fields = t.getAll();
	Object[] expected = results.get((String)fields.get(0));
	int i = 0;
	for (Object field : fields) {
	assertEquals(expected[i++], field);
	}
	}
	Util.deleteFile(cluster, "distinctAggs1Input.txt");

	}

	@Test
	public void testGroupElements() throws Exception {
	// test use of combiner when group elements are accessed in the foreach
	String input[] = {
	"ABC\t1\ta\t1",
	"ABC\t1\tb\t2",
	"ABC\t1\ta\t3",
	"ABC\t2\tb\t4",
	"DEF\t1\td\t1",
	"XYZ\t1\tx\t2"
	};

	Util.createInputFile(cluster, "testGroupElements.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'testGroupElements.txt' as (str:chararray, num1:int, alph : chararray, num2 : int);");
	pigServer.registerQuery("b = group a by (str, num1);");

	//check if combiner is present or not for various forms of foreach
	pigServer.registerQuery("c = foreach b generate flatten(group), COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", true);

	pigServer.registerQuery("c = foreach b generate group, COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", true);

	// projecting bag - combiner should not be used
	pigServer.registerQuery("c = foreach b generate group, a, COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", false);

	// projecting bag - combiner should not be used
	pigServer.registerQuery("c = foreach b generate group, a.num2, COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", false);

	pigServer.registerQuery("c = foreach b generate group.$0, group.$1, COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", true);

	pigServer.registerQuery("c = foreach b generate group.$0, group.$1 + COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", true);

	pigServer.registerQuery("c = foreach b generate group.str, group.$1, COUNT(a.alph), SUM(a.num2); ");
	checkCombinerUsed(pigServer, "c", true);

	pigServer.registerQuery("c = foreach b generate group.str, group.$1, COUNT(a.alph), SUM(a.num2), " +
	" (group.num1 == 1 ? (COUNT(a.num2) + 1) : (SUM(a.num2) + 10)) ; ");
	checkCombinerUsed(pigServer, "c", true);

	List<Tuple> expectedRes =
	Util.getTuplesFromConstantTupleStrings(
	new String[] {
	"('ABC',1,3L,6L,4L)",
	"('ABC',2,1L,4L,14L)",
	"('DEF',1,1L,1L,2L)",
	"('XYZ',1,1L,2L,2L)",
	});

	Iterator<Tuple> it = pigServer.openIterator("c");
	Util.checkQueryOutputsAfterSort(it, expectedRes);

	Util.deleteFile(cluster, "distinctAggs1Input.txt");

	}

	@Test
	public void testGroupByLimit() throws Exception {
	// test use of combiner when group elements are accessed in the foreach
	String input[] = {
	"ABC 1",
	"ABC 2",
	"DEF 1",
	"XYZ 1",
	"XYZ 2",
	"XYZ 3",
	};

	Util.createInputFile(cluster, "testGroupLimit.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'testGroupLimit.txt' using PigStorage(' ') " +
	"as (str:chararray, num1:int) ;");
	pigServer.registerQuery("b = group a by str;");


	pigServer.registerQuery("c = foreach b generate group, COUNT(a.num1) ; ");

	//check if combiner is present
	pigServer.registerQuery("d = limit c 2 ; ");
	checkCombinerUsed(pigServer, "d", true);

	List<Tuple> expectedRes =
	Util.getTuplesFromConstantTupleStrings(
	new String[] {
	"('ABC',2L)",
	"('DEF',1L)",
	});

	Iterator<Tuple> it = pigServer.openIterator("d");
	Util.checkQueryOutputsAfterSort(it, expectedRes);


	}

	private void checkCombinerUsed(PigServer pigServer, String string, boolean combineExpected)
	throws IOException {
	// make sure there is a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	boolean combinerFound = baos.toString().matches("(?si).combine plan.");
	System.out.println(baos.toString());
	assertEquals("is combiner present as expected", combineExpected, combinerFound);
	}


	@Test
	public void testDistinctNoCombiner() throws Exception {
	// test that combiner is NOT invoked when
	// one of the elements in the foreach generate
	// is a distinct() as the leaf
	String input[] = {
	"pig1\t18\t2.1",
	"pig2\t24\t3.3",
	"pig5\t45\t2.4",
	"pig1\t18\t2.1",
	"pig1\t19\t2.1",
	"pig2\t24\t4.5",
	"pig1\t20\t3.1" };

	Util.createInputFile(cluster, "distinctNoCombinerInput.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'distinctNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
	pigServer.registerQuery("b = group a by name;");
	pigServer.registerQuery("c = foreach b {" +
	" z = distinct a;" +
	" generate group, z, SUM(a.age), SUM(a.gpa);};");

	// make sure there is a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertFalse(baos.toString().matches("(?si).combine plan."));

	HashMap<String, Object[]> results = new HashMap<String, Object[]>();
	results.put("pig1", new Object[] {"pig1","bag-place-holder",75L,9.4});
	results.put("pig2", new Object[] {"pig2","bag-place-holder",48L,7.8});
	results.put("pig5", new Object[] {"pig5","bag-place-holder",45L,2.4});
	Iterator<Tuple> it = pigServer.openIterator("c");
	while(it.hasNext()) {
	Tuple t = it.next();
	List<Object> fields = t.getAll();
	Object[] expected = results.get((String)fields.get(0));
	int i = 0;
	for (Object field : fields) {
	if(i == 1) {
	// ignore the second field which is a bag
	// for comparison here
	continue;
	}
	assertEquals(expected[i++], field);
	}
	}
	Util.deleteFile(cluster, "distinctNoCombinerInput.txt");

	}

	@Test
	public void testForEachNoCombiner() throws Exception {
	// test that combiner is NOT invoked when
	// one of the elements in the foreach generate
	// has a foreach in the plan without a distinct agg
	String input[] = {
	"pig1\t18\t2.1",
	"pig2\t24\t3.3",
	"pig5\t45\t2.4",
	"pig1\t18\t2.1",
	"pig1\t19\t2.1",
	"pig2\t24\t4.5",
	"pig1\t20\t3.1" };

	Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
	pigServer.registerQuery("b = group a by name;");
	pigServer.registerQuery("c = foreach b {" +
	" z = a.age;" +
	" generate group, z, SUM(a.age), SUM(a.gpa);};");

	// make sure there is a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertFalse(baos.toString().matches("(?si).combine plan."));

	HashMap<String, Object[]> results = new HashMap<String, Object[]>();
	results.put("pig1", new Object[] {"pig1","bag-place-holder",75L,9.4});
	results.put("pig2", new Object[] {"pig2","bag-place-holder",48L,7.8});
	results.put("pig5", new Object[] {"pig5","bag-place-holder",45L,2.4});
	Iterator<Tuple> it = pigServer.openIterator("c");
	while(it.hasNext()) {
	Tuple t = it.next();
	List<Object> fields = t.getAll();
	Object[] expected = results.get((String)fields.get(0));
	int i = 0;
	for (Object field : fields) {
	if(i == 1) {
	// ignore the second field which is a bag
	// for comparison here
	continue;
	}
	assertEquals(expected[i++], field);
	}
	}
	Util.deleteFile(cluster, "forEachNoCombinerInput.txt");

	}

	@Test
	public void testJiraPig746() {
	// test that combiner is NOT invoked when
	// one of the elements in the foreach generate
	// has a foreach in the plan without a distinct agg
	String input[] = {
	"pig1\t18\t2.1",
	"pig2\t24\t3.3",
	"pig5\t45\t2.4",
	"pig1\t18\t2.1",
	"pig1\t19\t2.1",
	"pig2\t24\t4.5",
	"pig1\t20\t3.1" };

	String expected[] = {
	"(pig1,75,{(pig1,18,2.1),(pig1,18,2.1),(pig1,19,2.1),(pig1,20,3.1)})",
	"(pig2,48,{(pig2,24,3.3),(pig2,24,4.5)})",
	"(pig5,45,{(pig5,45,2.4)})"
	};

	try {
	Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);

	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
	pigServer.registerQuery("b = group a by name;");
	pigServer.registerQuery("c = foreach b generate group, SUM(a.age), a;");

	// make sure there isn't a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertFalse(baos.toString().matches("(?si).combine plan."));

	Iterator<Tuple> it = pigServer.openIterator("c");
	Util.checkQueryOutputsAfterSortRecursive(it, expected, "group:chararray,age:long,b:{t:(name:chararray,age:int,gpa:double)}");
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	} finally {
	try {
	Util.deleteFile(cluster, "forEachNoCombinerInput.txt");
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}

	public static class JiraPig1030 extends EvalFunc<DataBag> {

	public DataBag exec(Tuple input) throws IOException {
	return new DefaultDataBag();
	}
	}

	@Test
	public void testJiraPig1030() {
	// test that combiner is NOT invoked when
	// one of the elements in the foreach generate
	// has a non-algebraic UDF that have multiple inputs
	// (one of them is distinct).

	String input[] = {
	"pig1\t18\t2.1",
	"pig2\t24\t3.3",
	"pig5\t45\t2.4",
	"pig1\t18\t2.1",
	"pig1\t19\t2.1",
	"pig2\t24\t4.5",
	"pig1\t20\t3.1" };

	try {
	Util.createInputFile(cluster, "forEachNoCombinerInput.txt", input);
	PigServer pigServer = new PigServer(ExecType.MAPREDUCE, cluster.getProperties());
	pigServer.registerQuery("a = load 'forEachNoCombinerInput.txt' as (name:chararray, age:int, gpa:double);");
	pigServer.registerQuery("b = group a all;");
	pigServer.registerQuery("c = foreach b {" +
	" d = distinct a.age;" +
	" generate group, " + JiraPig1030.class.getName() + "(d, 0);};");

	// make sure there isn't a combine plan in the explain output
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	PrintStream ps = new PrintStream(baos);
	pigServer.explain("c", ps);
	assertFalse(baos.toString().matches("(?si).combine plan."));
	} catch (Exception e) {
	e.printStackTrace();
	Assert.fail();
	} finally {
	try {
	Util.deleteFile(cluster, "forEachNoCombinerInput.txt");
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}

	}