test/org/apache/pig/test/TestForEachNestedPlan.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 java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.text.DecimalFormat;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Random;

 import org.apache.pig.PigServer;
 import org.apache.pig.data.BagFactory;
 import org.apache.pig.data.DataBag;
 import org.apache.pig.data.Tuple;
 import org.apache.pig.parser.ParserException;
 import org.apache.pig.test.utils.TestHelper;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Test;

 public class TestForEachNestedPlan {

     static MiniGenericCluster cluster = MiniGenericCluster.buildCluster();

     private PigServer pig ;

     public TestForEachNestedPlan() throws Throwable {
         pig = new PigServer(cluster.getExecType(), cluster.getProperties()) ;
     }

     Boolean[] nullFlags = new Boolean[]{ false, true };

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

     @Test
     public void testInnerOrderBy() throws Exception {
         for (int i = 0; i < nullFlags.length; i++) {
             System.out.println("Running testInnerOrderBy with nullFlags set to :"
                             + nullFlags[i]);
             File tmpFile = genDataSetFile1(nullFlags[i]);
             pig.registerQuery("a = load '"
                     + Util.generateURI(tmpFile.toString(), pig.getPigContext()) + "'; ");
             pig.registerQuery("b = group a by $0; ");
             pig.registerQuery("c = foreach b { " + "     c1 = order $1 by *; "
                     + "    generate flatten(c1); " + "};");
             Iterator<Tuple> it = pig.openIterator("c");
             Tuple t = null;
             int count = 0;
             while (it.hasNext()) {
                 t = it.next();
                 System.out.println(count + ":" + t);
                 count++;
             }
             Assert.assertEquals(count, 30);
         }
     }

     @Test
     public void testInnerOrderByStarWithSchema() throws Exception {
         File tmpFile = genDataSetFile1(false);
         pig.registerQuery("a = load '" + Util.generateURI(tmpFile.toString(),
                 pig.getPigContext()) + "' as (a0, a1);");
         pig.registerQuery("b = group a by a0; ");
         pig.registerQuery("c = foreach b { d = order a by *; "
                 + "  generate group, d; };");
         Iterator<Tuple> it = pig.openIterator("c");
         Tuple t = null;
         int count = 0;
         while (it.hasNext()) {
             t = it.next();
             System.out.println(count + ":" + t);
             count++;
         }
         Assert.assertEquals(count, 10);
     }

     @Test
     public void testMultiColInAlias() throws Exception {
     	pig.getPigContext().getProperties().setProperty("pig.exec.nosecondarykey", "true");
     	String INPUT_FILE = "test-multi-alias.txt";
         PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
         w.println("10\tnrai01\t01");
         w.println("20\tnrai02\t02");
         w.close();

         try {

             Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);
             pig.registerQuery("A = load '" + INPUT_FILE + "' "
                     + "as (a:int, b:chararray, c:int);");
             pig.registerQuery("B = GROUP A BY (a, b);") ;

             DataBag dbfrj = BagFactory.getInstance().newDefaultBag(), dbshj = BagFactory.getInstance().newDefaultBag();
             {
                 pig.registerQuery("C = FOREACH B { bg = A.($1,$2); GENERATE group, bg; } ;") ;
                 Iterator<Tuple> iter1 = pig.openIterator("C");
                 while(iter1.hasNext()) {
                     dbfrj.add(iter1.next());
                 }
             }
             {
                 pig.registerQuery("D = FOREACH B { GENERATE group, A.($1,$2);};") ;
                 Iterator<Tuple> iter2 = pig.openIterator("D");
                 while(iter2.hasNext()) {
                     dbshj.add(iter2.next());
                 }
             }
             Assert.assertEquals(dbfrj.size(), dbshj.size());
             Assert.assertEquals(true, TestHelper.compareBags(dbfrj, dbshj));

         } finally{
             new File(INPUT_FILE).delete();
             try {
                 Util.deleteFile(cluster, INPUT_FILE);
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }

     @Test
     public void testAlgebricFuncWithoutGroupBy()
     throws IOException, ParserException {
         String INPUT_FILE = "test-sum.txt";

         PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
         w.println("10\t{(1),(2),(3)}");
         w.println("20\t{(4),(5),(6),(7)}");
         w.println("30\t{(8),(9)}");
         w.close();

         try {

             Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

             pig.registerQuery("a = load '" + INPUT_FILE + "' "
                     + "as (id:int, g:bag{t:tuple(u:int)});");
             pig.registerQuery("b = foreach a generate id, SUM(g);") ;

             Iterator<Tuple> iter = pig.openIterator("b");

             List<Tuple> expectedResults =
                 Util.getTuplesFromConstantTupleStrings(
                     new String[] {
                             "(10,6L)",
                             "(20,22L)",
                             "(30,17L)"
                     });

             int counter = 0;
             while (iter.hasNext()) {
                 assertEquals(expectedResults.get(counter++).toString(),
                         iter.next().toString());
             }

             assertEquals(expectedResults.size(), counter);

         } finally{
             new File(INPUT_FILE).delete();
             try {
                 Util.deleteFile(cluster, INPUT_FILE);
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }

     @Test
     public void testInnerDistinct()
     throws IOException, ParserException {
         String INPUT_FILE = "test-distinct.txt";

         PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
         w.println("10\t89");
         w.println("20\t78");
         w.println("10\t68");
         w.println("10\t89");
         w.println("20\t92");
         w.close();

         try {
             Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

             pig.registerQuery("A = load '" + INPUT_FILE
                     + "' as (age:int, gpa:int);");
             pig.registerQuery("B = group A by age;");
             pig.registerQuery("C = foreach B { D = A.gpa; E = distinct D; " +
             		"generate group, MIN(E); };");

             Iterator<Tuple> iter = pig.openIterator("C");
             List<Tuple> expectedResults =
                 Util.getTuplesFromConstantTupleStrings(
                         new String[] {"(10,68)", "(20,78)"});
             Util.checkQueryOutputsAfterSort(iter, expectedResults);
         } finally{
             new File(INPUT_FILE).delete();
             try {
                 Util.deleteFile(cluster, INPUT_FILE);
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }

     @Test
     public void testInnerOrderByAliasReuse()
     throws IOException, ParserException {
         String INPUT_FILE = "test-innerorderbyaliasreuse.txt";

         PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
         w.println("1\t4");
         w.println("1\t3");
         w.println("2\t3");
         w.println("2\t4");
         w.close();

         try {
             Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

             pig.registerQuery("A = load '" + INPUT_FILE
                     + "' as (v1:int, v2:int);");
             pig.registerQuery("B = group A by v1;");
             pig.registerQuery("C = foreach B { X = A; X = order X by v2 asc; " +
             		"generate flatten(X); };");

             Iterator<Tuple> iter = pig.openIterator("C");

             List<Tuple> expectedResults =
                 Util.getTuplesFromConstantTupleStrings(
                         new String[] {"(1,3)", "(1,4)", "(2,3)", "(2,4)"});
             Util.checkQueryOutputsAfterSort(iter, expectedResults);
         } finally{
             new File(INPUT_FILE).delete();
             try {
                 Util.deleteFile(cluster, INPUT_FILE);
             } catch (IOException e) {
                 e.printStackTrace();
                 Assert.fail();
             }
         }
     }


     /***
      * For generating a sample dataset
      */
     private File genDataSetFile1(boolean withNulls) throws IOException {

         int dataLength = 30;
         String[][] data = new String[dataLength][] ;

         DecimalFormat formatter = new DecimalFormat("0000000");

         Random r = new Random();

         for (int i = 0; i < dataLength; i++) {
             data[i] = new String[2] ;
             // inject nulls randomly
             if(withNulls && r.nextInt(dataLength) < 0.3 * dataLength) {
                 data[i][0] = "";
             } else {
                 data[i][0] = formatter.format(i % 10);
             }
             data[i][1] = formatter.format((dataLength - i)/2);
         }

         return TestHelper.createTempFile(data) ;
     }
 }
	/*
	* 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 java.io.File;
	import java.io.FileWriter;
	import java.io.IOException;
	import java.io.PrintWriter;
	import java.text.DecimalFormat;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Random;

	import org.apache.pig.PigServer;
	import org.apache.pig.data.BagFactory;
	import org.apache.pig.data.DataBag;
	import org.apache.pig.data.Tuple;
	import org.apache.pig.parser.ParserException;
	import org.apache.pig.test.utils.TestHelper;
	import org.junit.AfterClass;
	import org.junit.Assert;
	import org.junit.Test;

	public class TestForEachNestedPlan {

	static MiniGenericCluster cluster = MiniGenericCluster.buildCluster();

	private PigServer pig ;

	public TestForEachNestedPlan() throws Throwable {
	pig = new PigServer(cluster.getExecType(), cluster.getProperties()) ;
	}

	Boolean[] nullFlags = new Boolean[]{ false, true };

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

	@Test
	public void testInnerOrderBy() throws Exception {
	for (int i = 0; i < nullFlags.length; i++) {
	System.out.println("Running testInnerOrderBy with nullFlags set to :"
	+ nullFlags[i]);
	File tmpFile = genDataSetFile1(nullFlags[i]);
	pig.registerQuery("a = load '"
	+ Util.generateURI(tmpFile.toString(), pig.getPigContext()) + "'; ");
	pig.registerQuery("b = group a by $0; ");
	pig.registerQuery("c = foreach b { " + " c1 = order $1 by *; "
	+ " generate flatten(c1); " + "};");
	Iterator<Tuple> it = pig.openIterator("c");
	Tuple t = null;
	int count = 0;
	while (it.hasNext()) {
	t = it.next();
	System.out.println(count + ":" + t);
	count++;
	}
	Assert.assertEquals(count, 30);
	}
	}

	@Test
	public void testInnerOrderByStarWithSchema() throws Exception {
	File tmpFile = genDataSetFile1(false);
	pig.registerQuery("a = load '" + Util.generateURI(tmpFile.toString(),
	pig.getPigContext()) + "' as (a0, a1);");
	pig.registerQuery("b = group a by a0; ");
	pig.registerQuery("c = foreach b { d = order a by *; "
	+ " generate group, d; };");
	Iterator<Tuple> it = pig.openIterator("c");
	Tuple t = null;
	int count = 0;
	while (it.hasNext()) {
	t = it.next();
	System.out.println(count + ":" + t);
	count++;
	}
	Assert.assertEquals(count, 10);
	}

	@Test
	public void testMultiColInAlias() throws Exception {
	pig.getPigContext().getProperties().setProperty("pig.exec.nosecondarykey", "true");
	String INPUT_FILE = "test-multi-alias.txt";
	PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
	w.println("10\tnrai01\t01");
	w.println("20\tnrai02\t02");
	w.close();

	try {

	Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);
	pig.registerQuery("A = load '" + INPUT_FILE + "' "
	+ "as (a:int, b:chararray, c:int);");
	pig.registerQuery("B = GROUP A BY (a, b);") ;

	DataBag dbfrj = BagFactory.getInstance().newDefaultBag(), dbshj = BagFactory.getInstance().newDefaultBag();
	{
	pig.registerQuery("C = FOREACH B { bg = A.($1,$2); GENERATE group, bg; } ;") ;
	Iterator<Tuple> iter1 = pig.openIterator("C");
	while(iter1.hasNext()) {
	dbfrj.add(iter1.next());
	}
	}
	{
	pig.registerQuery("D = FOREACH B { GENERATE group, A.($1,$2);};") ;
	Iterator<Tuple> iter2 = pig.openIterator("D");
	while(iter2.hasNext()) {
	dbshj.add(iter2.next());
	}
	}
	Assert.assertEquals(dbfrj.size(), dbshj.size());
	Assert.assertEquals(true, TestHelper.compareBags(dbfrj, dbshj));

	} finally{
	new File(INPUT_FILE).delete();
	try {
	Util.deleteFile(cluster, INPUT_FILE);
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}

	@Test
	public void testAlgebricFuncWithoutGroupBy()
	throws IOException, ParserException {
	String INPUT_FILE = "test-sum.txt";

	PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
	w.println("10\t{(1),(2),(3)}");
	w.println("20\t{(4),(5),(6),(7)}");
	w.println("30\t{(8),(9)}");
	w.close();

	try {

	Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

	pig.registerQuery("a = load '" + INPUT_FILE + "' "
	+ "as (id:int, g:bag{t:tuple(u:int)});");
	pig.registerQuery("b = foreach a generate id, SUM(g);") ;

	Iterator<Tuple> iter = pig.openIterator("b");

	List<Tuple> expectedResults =
	Util.getTuplesFromConstantTupleStrings(
	new String[] {
	"(10,6L)",
	"(20,22L)",
	"(30,17L)"
	});

	int counter = 0;
	while (iter.hasNext()) {
	assertEquals(expectedResults.get(counter++).toString(),
	iter.next().toString());
	}

	assertEquals(expectedResults.size(), counter);

	} finally{
	new File(INPUT_FILE).delete();
	try {
	Util.deleteFile(cluster, INPUT_FILE);
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}

	@Test
	public void testInnerDistinct()
	throws IOException, ParserException {
	String INPUT_FILE = "test-distinct.txt";

	PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
	w.println("10\t89");
	w.println("20\t78");
	w.println("10\t68");
	w.println("10\t89");
	w.println("20\t92");
	w.close();

	try {
	Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

	pig.registerQuery("A = load '" + INPUT_FILE
	+ "' as (age:int, gpa:int);");
	pig.registerQuery("B = group A by age;");
	pig.registerQuery("C = foreach B { D = A.gpa; E = distinct D; " +
	"generate group, MIN(E); };");

	Iterator<Tuple> iter = pig.openIterator("C");
	List<Tuple> expectedResults =
	Util.getTuplesFromConstantTupleStrings(
	new String[] {"(10,68)", "(20,78)"});
	Util.checkQueryOutputsAfterSort(iter, expectedResults);
	} finally{
	new File(INPUT_FILE).delete();
	try {
	Util.deleteFile(cluster, INPUT_FILE);
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}

	@Test
	public void testInnerOrderByAliasReuse()
	throws IOException, ParserException {
	String INPUT_FILE = "test-innerorderbyaliasreuse.txt";

	PrintWriter w = new PrintWriter(new FileWriter(INPUT_FILE));
	w.println("1\t4");
	w.println("1\t3");
	w.println("2\t3");
	w.println("2\t4");
	w.close();

	try {
	Util.copyFromLocalToCluster(cluster, INPUT_FILE, INPUT_FILE);

	pig.registerQuery("A = load '" + INPUT_FILE
	+ "' as (v1:int, v2:int);");
	pig.registerQuery("B = group A by v1;");
	pig.registerQuery("C = foreach B { X = A; X = order X by v2 asc; " +
	"generate flatten(X); };");

	Iterator<Tuple> iter = pig.openIterator("C");

	List<Tuple> expectedResults =
	Util.getTuplesFromConstantTupleStrings(
	new String[] {"(1,3)", "(1,4)", "(2,3)", "(2,4)"});
	Util.checkQueryOutputsAfterSort(iter, expectedResults);
	} finally{
	new File(INPUT_FILE).delete();
	try {
	Util.deleteFile(cluster, INPUT_FILE);
	} catch (IOException e) {
	e.printStackTrace();
	Assert.fail();
	}
	}
	}


	/***
	* For generating a sample dataset
	*/
	private File genDataSetFile1(boolean withNulls) throws IOException {

	int dataLength = 30;
	String[][] data = new String[dataLength][] ;

	DecimalFormat formatter = new DecimalFormat("0000000");

	Random r = new Random();

	for (int i = 0; i < dataLength; i++) {
	data[i] = new String[2] ;
	// inject nulls randomly
	if(withNulls && r.nextInt(dataLength) < 0.3 * dataLength) {
	data[i][0] = "";
	} else {
	data[i][0] = formatter.format(i % 10);
	}
	data[i][1] = formatter.format((dataLength - i)/2);
	}

	return TestHelper.createTempFile(data) ;
	}
	}