test/org/apache/pig/test/TestCustomPartitioner.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 java.io.BufferedReader;
 import java.io.File;
 import java.io.FileReader;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Properties;
 import java.util.Random;

 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.pig.PigServer;
 import org.apache.pig.data.BagFactory;
 import org.apache.pig.data.Tuple;
 import org.apache.pig.data.TupleFactory;
 import org.apache.pig.impl.io.FileLocalizer;
 import org.apache.pig.impl.logicalLayer.FrontendException;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;

 public class TestCustomPartitioner {
     private static MiniGenericCluster cluster;
     private static Properties properties;
     private static PigServer pigServer;
     private static FileSystem fs;

     TupleFactory mTf = TupleFactory.getInstance();
     BagFactory mBf = BagFactory.getInstance();

     @Before
     public void setUp() throws Exception {
         FileLocalizer.setR(new Random());
         pigServer = new PigServer(cluster.getExecType(), properties);
     }

     @BeforeClass
     public static void oneTimeSetUp() throws Exception {
         cluster = MiniGenericCluster.buildCluster();
         properties = cluster.getProperties();
         fs = cluster.getFileSystem();
     }

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

     // See PIG-282
     @Test
     public void testCustomPartitionerParseJoins() throws Exception{
         String[] input = {
                 "1\t3",
                 "1\t2"
         };
         Util.createInputFile(cluster, "table_testCustomPartitionerParseJoins", input);

         // Custom Partitioner is not allowed for skewed joins, will throw a ExecException
         try {
             pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerParseJoins' as (a0:int, a1:int);");
             pigServer.registerQuery("B = ORDER A by $0;");
             pigServer.registerQuery("skewed = JOIN A by $0, B by $0 USING 'skewed' PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
             //control should not reach here
             Assert.fail("Skewed join cannot accept a custom partitioner");
         } catch(FrontendException e) {
             Assert.assertTrue( e.getMessage().contains( "Custom Partitioner is not supported for skewed join" ) );
         }

         pigServer.registerQuery("hash = JOIN A by $0, B by $0 USING 'hash' PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
         Iterator<Tuple> iter = pigServer.openIterator("hash");

         List<String> expected = new ArrayList<String>();
         expected.add("(1,3,1,2)");
         expected.add("(1,3,1,3)");
         expected.add("(1,2,1,2)");
         expected.add("(1,2,1,3)");
         Collections.sort(expected);

         List<String> actual = new ArrayList<String>();
         while (iter.hasNext()) {
             actual.add(iter.next().toString());
         }
         Collections.sort(actual);

         Assert.assertEquals(expected, actual);

         // No checks are made for merged and replicated joins as they are compiled to a map only job
         // No frontend error checking has been added for these jobs, hence not adding any test cases
         // Manually tested the sanity once. Above test should cover the basic sanity of the scenario

         Util.deleteFile(cluster, "table_testCustomPartitionerParseJoins");
     }

     // See PIG-282
     @Test
     public void testCustomPartitionerGroups() throws Exception{
         String[] input = {
                 "1\t1",
                 "2\t1",
                 "3\t1",
                 "4\t1"
         };
         Util.createInputFile(cluster, "table_testCustomPartitionerGroups", input);

         String outputDir = "tmp_testCustomPartitionerGroup";
         pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerGroups' as (a0:int, a1:int);");
         // It should be noted that for a map reduce job, the total number of partitions
         // is the same as the number of reduce tasks for the job. Hence we need to find a case wherein
         // we will get more than one reduce job so that we can use the partitioner.
         // The following logic assumes that we get 2 reduce jobs, so that we can hard-code the logic.
         // SimpleCustomPartitioner3 simply returns '1' (second reducer) for all inputs when
         // partition number is bigger than 1.
         //
         pigServer.registerQuery("B = group A by $0 PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner3 parallel 2;");
         pigServer.store("B", outputDir);

         new File(outputDir).mkdir();
         FileStatus[] outputFiles = fs.listStatus(new Path(outputDir), Util.getSuccessMarkerPathFilter());

         Util.copyFromClusterToLocal(cluster, outputFiles[0].getPath().toString(), outputDir + "/" + 0);
         BufferedReader reader = new BufferedReader(new FileReader(outputDir + "/" + 0));
         while(reader.readLine() != null) {
             Assert.fail("Partition 0 should be empty.  Most likely Custom Partitioner was not used.");
         }
         reader.close();

         Util.copyFromClusterToLocal(cluster, outputFiles[1].getPath().toString(), outputDir + "/" + 1);
         reader = new BufferedReader(new FileReader(outputDir + "/" + 1));
         int count=0;
         while(reader.readLine() != null) {
             //all outputs should come to partion 1 (with SimpleCustomPartitioner3)
             count++;
         }
         reader.close();
         Assert.assertEquals(4, count);

         Util.deleteDirectory(new File(outputDir));
         Util.deleteFile(cluster, outputDir);
         Util.deleteFile(cluster, "table_testCustomPartitionerGroups");
     }

     // See PIG-3385
     @Test
     public void testCustomPartitionerDistinct() throws Exception{
         String[] input = {
                 "1\t1",
                 "2\t1",
                 "1\t1",
                 "3\t1",
                 "4\t1",
         };
         Util.createInputFile(cluster, "table_testCustomPartitionerDistinct", input);

         String outputDir = "tmp_testCustomPartitionerDistinct";
         pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerDistinct' as (a0:int, a1:int);");
         pigServer.registerQuery("B = distinct A PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner3 parallel 2;");
         pigServer.store("B", outputDir);

         new File(outputDir).mkdir();
         FileStatus[] outputFiles = fs.listStatus(new Path(outputDir), Util.getSuccessMarkerPathFilter());

         // SimpleCustomPartitioner3 simply partition all inputs to *second* reducer
         Util.copyFromClusterToLocal(cluster, outputFiles[0].getPath().toString(), outputDir + "/" + 0);
         BufferedReader reader = new BufferedReader(new FileReader(outputDir + "/" + 0));
         while (reader.readLine() != null) {
             Assert.fail("Partition 0 should be empty.  Most likely Custom Partitioner was not used.");
         }
         reader.close();

         Util.copyFromClusterToLocal(cluster, outputFiles[1].getPath().toString(), outputDir + "/" + 1);
         reader = new BufferedReader(new FileReader(outputDir + "/" + 1));
         int count=0;
         while (reader.readLine() != null) {
             //all outputs should come to partion 1 (with SimpleCustomPartitioner3)
             count++;
         }
         reader.close();
         Assert.assertEquals(4, count);

         Util.deleteDirectory(new File(outputDir));
         Util.deleteFile(cluster, outputDir);
         Util.deleteFile(cluster, "table_testCustomPartitionerDistinct");
     }

     // See PIG-282
     @Test
     public void testCustomPartitionerCross() throws Exception{
     	String[] input = {
                 "1\t3",
                 "1\t2",
         };

         Util.createInputFile(cluster, "table_testCustomPartitionerCross", input);
         pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerCross' as (a0:int, a1:int);");
         pigServer.registerQuery("B = ORDER A by $0;");
         pigServer.registerQuery("C = cross A , B PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
         Iterator<Tuple> iter = pigServer.openIterator("C");
         Tuple t;

         Collection<String> results = new HashSet<String>();
         results.add("(1,3,1,2)");
         results.add("(1,3,1,3)");
         results.add("(1,2,1,2)");
         results.add("(1,2,1,3)");

         Assert.assertTrue(iter.hasNext());
         t = iter.next();
         Assert.assertTrue(t.size()==4);
         Assert.assertTrue(results.contains(t.toString()));

         Assert.assertTrue(iter.hasNext());
         t = iter.next();
         Assert.assertTrue(t.size()==4);
         Assert.assertTrue(results.contains(t.toString()));

         Assert.assertTrue(iter.hasNext());
         t = iter.next();
         Assert.assertTrue(t.size()==4);
         Assert.assertTrue(results.contains(t.toString()));

         Assert.assertTrue(iter.hasNext());
         t = iter.next();
         Assert.assertTrue(t.size()==4);
         Assert.assertTrue(results.contains(t.toString()));

         Util.deleteFile(cluster, "table_testCustomPartitionerCross");
     }
 }
	/*
	* 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 java.io.BufferedReader;
	import java.io.File;
	import java.io.FileReader;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.HashSet;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Properties;
	import java.util.Random;

	import org.apache.hadoop.fs.FileStatus;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.Path;
	import org.apache.pig.PigServer;
	import org.apache.pig.data.BagFactory;
	import org.apache.pig.data.Tuple;
	import org.apache.pig.data.TupleFactory;
	import org.apache.pig.impl.io.FileLocalizer;
	import org.apache.pig.impl.logicalLayer.FrontendException;
	import org.junit.AfterClass;
	import org.junit.Assert;
	import org.junit.Before;
	import org.junit.BeforeClass;
	import org.junit.Test;

	public class TestCustomPartitioner {
	private static MiniGenericCluster cluster;
	private static Properties properties;
	private static PigServer pigServer;
	private static FileSystem fs;

	TupleFactory mTf = TupleFactory.getInstance();
	BagFactory mBf = BagFactory.getInstance();

	@Before
	public void setUp() throws Exception {
	FileLocalizer.setR(new Random());
	pigServer = new PigServer(cluster.getExecType(), properties);
	}

	@BeforeClass
	public static void oneTimeSetUp() throws Exception {
	cluster = MiniGenericCluster.buildCluster();
	properties = cluster.getProperties();
	fs = cluster.getFileSystem();
	}

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

	// See PIG-282
	@Test
	public void testCustomPartitionerParseJoins() throws Exception{
	String[] input = {
	"1\t3",
	"1\t2"
	};
	Util.createInputFile(cluster, "table_testCustomPartitionerParseJoins", input);

	// Custom Partitioner is not allowed for skewed joins, will throw a ExecException
	try {
	pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerParseJoins' as (a0:int, a1:int);");
	pigServer.registerQuery("B = ORDER A by $0;");
	pigServer.registerQuery("skewed = JOIN A by $0, B by $0 USING 'skewed' PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
	//control should not reach here
	Assert.fail("Skewed join cannot accept a custom partitioner");
	} catch(FrontendException e) {
	Assert.assertTrue( e.getMessage().contains( "Custom Partitioner is not supported for skewed join" ) );
	}

	pigServer.registerQuery("hash = JOIN A by $0, B by $0 USING 'hash' PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
	Iterator<Tuple> iter = pigServer.openIterator("hash");

	List<String> expected = new ArrayList<String>();
	expected.add("(1,3,1,2)");
	expected.add("(1,3,1,3)");
	expected.add("(1,2,1,2)");
	expected.add("(1,2,1,3)");
	Collections.sort(expected);

	List<String> actual = new ArrayList<String>();
	while (iter.hasNext()) {
	actual.add(iter.next().toString());
	}
	Collections.sort(actual);

	Assert.assertEquals(expected, actual);

	// No checks are made for merged and replicated joins as they are compiled to a map only job
	// No frontend error checking has been added for these jobs, hence not adding any test cases
	// Manually tested the sanity once. Above test should cover the basic sanity of the scenario

	Util.deleteFile(cluster, "table_testCustomPartitionerParseJoins");
	}

	// See PIG-282
	@Test
	public void testCustomPartitionerGroups() throws Exception{
	String[] input = {
	"1\t1",
	"2\t1",
	"3\t1",
	"4\t1"
	};
	Util.createInputFile(cluster, "table_testCustomPartitionerGroups", input);

	String outputDir = "tmp_testCustomPartitionerGroup";
	pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerGroups' as (a0:int, a1:int);");
	// It should be noted that for a map reduce job, the total number of partitions
	// is the same as the number of reduce tasks for the job. Hence we need to find a case wherein
	// we will get more than one reduce job so that we can use the partitioner.
	// The following logic assumes that we get 2 reduce jobs, so that we can hard-code the logic.
	// SimpleCustomPartitioner3 simply returns '1' (second reducer) for all inputs when
	// partition number is bigger than 1.
	//
	pigServer.registerQuery("B = group A by $0 PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner3 parallel 2;");
	pigServer.store("B", outputDir);

	new File(outputDir).mkdir();
	FileStatus[] outputFiles = fs.listStatus(new Path(outputDir), Util.getSuccessMarkerPathFilter());

	Util.copyFromClusterToLocal(cluster, outputFiles[0].getPath().toString(), outputDir + "/" + 0);
	BufferedReader reader = new BufferedReader(new FileReader(outputDir + "/" + 0));
	while(reader.readLine() != null) {
	Assert.fail("Partition 0 should be empty. Most likely Custom Partitioner was not used.");
	}
	reader.close();

	Util.copyFromClusterToLocal(cluster, outputFiles[1].getPath().toString(), outputDir + "/" + 1);
	reader = new BufferedReader(new FileReader(outputDir + "/" + 1));
	int count=0;
	while(reader.readLine() != null) {
	//all outputs should come to partion 1 (with SimpleCustomPartitioner3)
	count++;
	}
	reader.close();
	Assert.assertEquals(4, count);

	Util.deleteDirectory(new File(outputDir));
	Util.deleteFile(cluster, outputDir);
	Util.deleteFile(cluster, "table_testCustomPartitionerGroups");
	}

	// See PIG-3385
	@Test
	public void testCustomPartitionerDistinct() throws Exception{
	String[] input = {
	"1\t1",
	"2\t1",
	"1\t1",
	"3\t1",
	"4\t1",
	};
	Util.createInputFile(cluster, "table_testCustomPartitionerDistinct", input);

	String outputDir = "tmp_testCustomPartitionerDistinct";
	pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerDistinct' as (a0:int, a1:int);");
	pigServer.registerQuery("B = distinct A PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner3 parallel 2;");
	pigServer.store("B", outputDir);

	new File(outputDir).mkdir();
	FileStatus[] outputFiles = fs.listStatus(new Path(outputDir), Util.getSuccessMarkerPathFilter());

	// SimpleCustomPartitioner3 simply partition all inputs to second reducer
	Util.copyFromClusterToLocal(cluster, outputFiles[0].getPath().toString(), outputDir + "/" + 0);
	BufferedReader reader = new BufferedReader(new FileReader(outputDir + "/" + 0));
	while (reader.readLine() != null) {
	Assert.fail("Partition 0 should be empty. Most likely Custom Partitioner was not used.");
	}
	reader.close();

	Util.copyFromClusterToLocal(cluster, outputFiles[1].getPath().toString(), outputDir + "/" + 1);
	reader = new BufferedReader(new FileReader(outputDir + "/" + 1));
	int count=0;
	while (reader.readLine() != null) {
	//all outputs should come to partion 1 (with SimpleCustomPartitioner3)
	count++;
	}
	reader.close();
	Assert.assertEquals(4, count);

	Util.deleteDirectory(new File(outputDir));
	Util.deleteFile(cluster, outputDir);
	Util.deleteFile(cluster, "table_testCustomPartitionerDistinct");
	}

	// See PIG-282
	@Test
	public void testCustomPartitionerCross() throws Exception{
	String[] input = {
	"1\t3",
	"1\t2",
	};

	Util.createInputFile(cluster, "table_testCustomPartitionerCross", input);
	pigServer.registerQuery("A = LOAD 'table_testCustomPartitionerCross' as (a0:int, a1:int);");
	pigServer.registerQuery("B = ORDER A by $0;");
	pigServer.registerQuery("C = cross A , B PARTITION BY org.apache.pig.test.utils.SimpleCustomPartitioner;");
	Iterator<Tuple> iter = pigServer.openIterator("C");
	Tuple t;

	Collection<String> results = new HashSet<String>();
	results.add("(1,3,1,2)");
	results.add("(1,3,1,3)");
	results.add("(1,2,1,2)");
	results.add("(1,2,1,3)");

	Assert.assertTrue(iter.hasNext());
	t = iter.next();
	Assert.assertTrue(t.size()==4);
	Assert.assertTrue(results.contains(t.toString()));

	Assert.assertTrue(iter.hasNext());
	t = iter.next();
	Assert.assertTrue(t.size()==4);
	Assert.assertTrue(results.contains(t.toString()));

	Assert.assertTrue(iter.hasNext());
	t = iter.next();
	Assert.assertTrue(t.size()==4);
	Assert.assertTrue(results.contains(t.toString()));

	Assert.assertTrue(iter.hasNext());
	t = iter.next();
	Assert.assertTrue(t.size()==4);
	Assert.assertTrue(results.contains(t.toString()));

	Util.deleteFile(cluster, "table_testCustomPartitionerCross");
	}
	}