tez-runtime-library/src/test/java/org/apache/tez/runtime/library/output/TestOnFileSortedOutput.java - tez - 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
  * <p/>
  * http://www.apache.org/licenses/LICENSE-2.0
  * <p/>
  * 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.tez.runtime.library.output;

 import com.google.protobuf.ByteString;

 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.DataOutputBuffer;
 import org.apache.hadoop.io.Text;
 import org.apache.tez.common.TezUtils;
 import org.apache.tez.common.TezRuntimeFrameworkConfigs;
 import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.dag.api.TezConfiguration;
 import org.apache.tez.dag.api.UserPayload;
 import org.apache.tez.runtime.api.ExecutionContext;
 import org.apache.tez.runtime.api.Event;
 import org.apache.tez.runtime.api.MemoryUpdateCallback;
 import org.apache.tez.runtime.api.OutputContext;
 import org.apache.tez.runtime.api.OutputStatisticsReporter;
 import org.apache.tez.runtime.api.events.CompositeDataMovementEvent;
 import org.apache.tez.runtime.api.events.VertexManagerEvent;
 import org.apache.tez.runtime.library.api.KeyValuesWriter;
 import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
 import org.apache.tez.runtime.library.api.TezRuntimeConfiguration.ReportPartitionStats;
 import org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler;
 import org.apache.tez.runtime.library.common.sort.impl.dflt.DefaultSorter;
 import org.apache.tez.runtime.library.conf.OrderedPartitionedKVOutputConfig.SorterImpl;
 import org.apache.tez.runtime.library.partitioner.HashPartitioner;
 import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;

 import java.io.IOException;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicLong;

 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyLong;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;

 @SuppressWarnings({ "rawtypes", "unchecked" })
 @RunWith(Parameterized.class)
 public class TestOnFileSortedOutput {
   private static final Random rnd = new Random();
   private static final String UniqueID = "UUID";
   private static final String HOST = "localhost";
   private static final int PORT = 80;

   private Configuration conf;
   private FileSystem fs;
   private Path workingDir;
   //no of outputs
   private int partitions;
   //For sorter (pipelined / Default)
   private SorterImpl sorterImpl;
   private int sorterThreads;

   final AtomicLong outputSize = new AtomicLong();
   final AtomicLong numRecords = new AtomicLong();

   private KeyValuesWriter writer;
   private OrderedPartitionedKVOutput sortedOutput;
   private boolean sendEmptyPartitionViaEvent;
   //Partition index for which data should not be written to.
   private int emptyPartitionIdx;
   private ReportPartitionStats reportPartitionStats;

   /**
    * Constructor
    *
    * @param sendEmptyPartitionViaEvent
    * @param sorterImpl Which sorter impl ( pipeline/legacy )
    * @param sorterThreads number of threads needed for sorter (required only for pipelined sorter)
    * @param emptyPartitionIdx for which data should not be generated
    */
   public TestOnFileSortedOutput(boolean sendEmptyPartitionViaEvent,
       SorterImpl sorterImpl, int sorterThreads, int emptyPartitionIdx,
       ReportPartitionStats reportPartitionStats) throws IOException {
     this.sendEmptyPartitionViaEvent = sendEmptyPartitionViaEvent;
     this.emptyPartitionIdx = emptyPartitionIdx;
     this.sorterImpl = sorterImpl;
     this.sorterThreads = sorterThreads;
     this.reportPartitionStats = reportPartitionStats;
     conf = new Configuration();

     workingDir = new Path(".", this.getClass().getName());
     String localDirs = workingDir.toString();
     conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, localDirs);
     fs = FileSystem.getLocal(conf);
   }

   @Before
   public void setup() throws Exception {
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, sorterImpl.name());
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SORTER_SORT_THREADS, sorterThreads);
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 5);

     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS, Text.class.getName());
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS, Text.class.getName());
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS,
         HashPartitioner.class.getName());

     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_EMPTY_PARTITION_INFO_VIA_EVENTS_ENABLED,
         sendEmptyPartitionViaEvent);
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS,
         reportPartitionStats.getType());
     outputSize.set(0);
     numRecords.set(0);
     fs.mkdirs(workingDir);
     this.partitions = Math.max(1, rnd.nextInt(10));
   }

   @After
   public void cleanup() throws IOException {
     fs.delete(workingDir, true);
   }

   @SuppressWarnings("deprecation")
   @Parameterized.Parameters(name = "test[{0}, {1}, {2}, {3}, {4}]")
   public static Collection<Object[]> getParameters() {
     Collection<Object[]> parameters = new ArrayList<Object[]>();
     //empty_partition_via_events_enabled, noOfSortThreads,
     // partitionToBeEmpty, reportPartitionStats
     parameters.add(new Object[] { false, SorterImpl.LEGACY, 1, -1,
         ReportPartitionStats.ENABLED });
     parameters.add(new Object[] { false, SorterImpl.LEGACY, 1, 0,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, -1,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, 0,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, 0,
         ReportPartitionStats.PRECISE  });

     //Pipelined sorter
     parameters.add(new Object[] { false, SorterImpl.PIPELINED, 2, -1,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { false, SorterImpl.PIPELINED, 2, 0,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, -1,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, 0,
         ReportPartitionStats.ENABLED  });
     parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, 0,
         ReportPartitionStats.PRECISE  });
     return parameters;
   }

   private void startSortedOutput(int partitions) throws Exception {
     OutputContext context = createTezOutputContext();
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 4);
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
     sortedOutput.initialize();
     sortedOutput.start();
     writer = sortedOutput.getWriter();
   }

   private void _testPipelinedShuffle(String sorterName) throws Exception {
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);

     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, sorterName);
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);
     OutputContext context = createTezOutputContext();
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

     sortedOutput.initialize();
     sortedOutput.start();

     assertFalse(sortedOutput.finalMergeEnabled);
     assertTrue(sortedOutput.pipelinedShuffle);
   }

   @Test (timeout = 5000)
   public void testPipelinedShuffle() throws Exception {
     _testPipelinedShuffle(SorterImpl.PIPELINED.name());
   }

   @Test (timeout = 5000)
   public void testPipelinedShuffleWithFinalMerge() throws Exception {
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, SorterImpl.PIPELINED.name());

     //wrong setting for final merge enable in output
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, true);
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);
     OutputContext context = createTezOutputContext();
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

     sortedOutput.initialize();
     sortedOutput.start();
     assertFalse(sortedOutput.finalMergeEnabled); //should be disabled as pipelining is on
     assertTrue(sortedOutput.pipelinedShuffle);
   }

   @Test
   public void testPipelinedSettingsWithDefaultSorter() throws Exception {
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);
     //negative. with default sorter
     conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, SorterImpl.LEGACY.name());
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
     conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);

     OutputContext context = createTezOutputContext();
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

     sortedOutput.initialize();
     try {
       sortedOutput.start();
       fail("Should have thrown illegal argument exception as pipelining is enabled with "
           + "DefaultSorter");
     } catch(IllegalArgumentException ie) {
       assertTrue(ie.getMessage().contains("works with PipelinedSorter"));
     }

   }

   @Test (timeout = 5000)
   public void testSortBufferSize() throws Exception{
     OutputContext context = createTezOutputContext();
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 2048);
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
     try {
       //Memory limit checks are done in sorter impls. For e.g, defaultsorter does not support > 2GB
       sortedOutput.initialize();
       DefaultSorter sorter = new DefaultSorter(context, conf, 100, 3500*1024*1024l);
       fail();
     } catch(IllegalArgumentException e) {
       assertTrue(e.getMessage().contains(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB));
     }
     conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 0);
     payLoad = TezUtils.createUserPayloadFromConf(conf);
     doReturn(payLoad).when(context).getUserPayload();
     sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
     try {
       sortedOutput.initialize();
       fail();
     } catch(IllegalArgumentException e) {
       assertTrue(e.getMessage().contains(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB));
     }
   }

   @Test(timeout = 5000)
   public void baseTest() throws Exception {
     startSortedOutput(partitions);

     //Write random set of keys
     long recordsWritten = numRecords.get();
     for (int i = 0; i < Math.max(1, rnd.nextInt(50)); i++) {
       Text key = new Text(new BigInteger(256, rnd).toString());
       LinkedList values = new LinkedList();
       for (int j = 0; j < Math.max(2, rnd.nextInt(10)); j++) {
         recordsWritten++;
         values.add(new Text(new BigInteger(256, rnd).toString()));
       }
       writer.write(key, values);
     }

     List<Event> eventList = sortedOutput.close();
     assertTrue(eventList != null && eventList.size() == 2);
     assertEquals(recordsWritten, numRecords.get());
     ShuffleUserPayloads.DataMovementEventPayloadProto
         payload = ShuffleUserPayloads.DataMovementEventPayloadProto
         .parseFrom(
             ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));

     ShuffleUserPayloads.VertexManagerEventPayloadProto
         vmPayload = ShuffleUserPayloads.VertexManagerEventPayloadProto
         .parseFrom(
             ByteString.copyFrom(((VertexManagerEvent) eventList.get(0)).getUserPayload()));

     if (reportPartitionStats.isPrecise()) {
       assertTrue(vmPayload.hasDetailedPartitionStats());
     } else {
       assertTrue(vmPayload.hasPartitionStats());
     }
     assertEquals(HOST, payload.getHost());
     assertEquals(PORT, payload.getPort());
     assertEquals(UniqueID, payload.getPathComponent());
   }

   @Test(timeout = 5000)
   public void testWithSomeEmptyPartition() throws Exception {
     //ensure atleast 2 partitions are available
     partitions = Math.max(2, partitions);
     startSortedOutput(partitions);

     //write random data
     for (int i = 0; i < 2 * partitions; i++) {
       Text key = new Text(new BigInteger(256, rnd).toString());
       Text value = new Text(new BigInteger(256, rnd).toString());
       //skip writing to certain partitions
       if (i % partitions != emptyPartitionIdx) {
         writer.write(key, value);
       }
     }

     List<Event> eventList = sortedOutput.close();
     assertTrue(eventList != null && eventList.size() == 2);

     ShuffleUserPayloads.DataMovementEventPayloadProto
         payload = ShuffleUserPayloads.DataMovementEventPayloadProto
         .parseFrom(ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));

     assertEquals(HOST, payload.getHost());
     assertEquals(PORT, payload.getPort());
     assertEquals(UniqueID, payload.getPathComponent());
   }

   @Test(timeout = 5000)
   public void testAllEmptyPartition() throws Exception {
     startSortedOutput(partitions);

     //Close output without writing any data to it.
     List<Event> eventList = sortedOutput.close();
     assertTrue(eventList != null && eventList.size() == 2);

     ShuffleUserPayloads.DataMovementEventPayloadProto
         payload = ShuffleUserPayloads.DataMovementEventPayloadProto
         .parseFrom(ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));
     if (sendEmptyPartitionViaEvent) {
       assertEquals("", payload.getHost());
       assertEquals(0, payload.getPort());
       assertEquals("", payload.getPathComponent());
     } else {
       assertEquals(HOST, payload.getHost());
       assertEquals(PORT, payload.getPort());
       assertEquals(UniqueID, payload.getPathComponent());
     }
   }

   private OutputContext createTezOutputContext() throws IOException {
     String[] workingDirs = { workingDir.toString() };
     Configuration localConf = new Configuration(false);
     UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
     DataOutputBuffer serviceProviderMetaData = new DataOutputBuffer();
     serviceProviderMetaData.writeInt(PORT);

     TezCounters counters = new TezCounters();

     OutputStatisticsReporter reporter = mock(OutputStatisticsReporter.class);
     doAnswer(new Answer() {
       @Override public Object answer(InvocationOnMock invocation) throws Throwable {
         outputSize.set((Long) invocation.getArguments()[0]);
         return null;
       }
     }).when(reporter).reportDataSize(anyLong());
     doAnswer(new Answer() {
       @Override public Object answer(InvocationOnMock invocation) throws Throwable {
         numRecords.set((Long) invocation.getArguments()[0]);
         return null;
       }
     }).when(reporter).reportItemsProcessed(anyLong());


     OutputContext context = mock(OutputContext.class);
     doReturn(localConf).when(context).getContainerConfiguration();
     doReturn(counters).when(context).getCounters();
     doReturn(workingDirs).when(context).getWorkDirs();
     doReturn(payLoad).when(context).getUserPayload();
     doReturn(5 * 1024 * 1024l).when(context).getTotalMemoryAvailableToTask();
     doReturn(UniqueID).when(context).getUniqueIdentifier();
     doReturn("v0").when(context).getTaskVertexName();
     doReturn("v1").when(context).getDestinationVertexName();
     doReturn(ByteBuffer.wrap(serviceProviderMetaData.getData())).when(context)
         .getServiceProviderMetaData
             (conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID,
                 TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
     doAnswer(new Answer() {
       @Override public Object answer(InvocationOnMock invocation) throws Throwable {
         long requestedSize = (Long) invocation.getArguments()[0];
         MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation
             .getArguments()[1];
         callback.memoryAssigned(requestedSize);
         return null;
       }
     }).when(context).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
     ExecutionContext ExecutionContext = mock(ExecutionContext.class);
     doReturn(HOST).when(ExecutionContext).getHostName();
     doReturn(reporter).when(context).getStatisticsReporter();
     doReturn(ExecutionContext).when(context).getExecutionContext();
     return context;
   }

   @Test(timeout=5000)
   public void testInvalidSorter() throws Exception {
     try {
       _testPipelinedShuffle("Foo");
       Assert.fail("Expected start to fail due to invalid sorter");
     } catch (IllegalArgumentException e) {
       // Expected
     }
   }

   @Test(timeout=5000)
   public void testLowerCaseNamedSorter() throws Exception {
     _testPipelinedShuffle("Pipelined");
   }


 }
	/**
	* 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
	* <p/>
	* http://www.apache.org/licenses/LICENSE-2.0
	* <p/>
	* 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.tez.runtime.library.output;

	import com.google.protobuf.ByteString;

	import org.apache.hadoop.conf.Configuration;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.Path;
	import org.apache.hadoop.io.DataOutputBuffer;
	import org.apache.hadoop.io.Text;
	import org.apache.tez.common.TezUtils;
	import org.apache.tez.common.TezRuntimeFrameworkConfigs;
	import org.apache.tez.common.counters.TezCounters;
	import org.apache.tez.dag.api.TezConfiguration;
	import org.apache.tez.dag.api.UserPayload;
	import org.apache.tez.runtime.api.ExecutionContext;
	import org.apache.tez.runtime.api.Event;
	import org.apache.tez.runtime.api.MemoryUpdateCallback;
	import org.apache.tez.runtime.api.OutputContext;
	import org.apache.tez.runtime.api.OutputStatisticsReporter;
	import org.apache.tez.runtime.api.events.CompositeDataMovementEvent;
	import org.apache.tez.runtime.api.events.VertexManagerEvent;
	import org.apache.tez.runtime.library.api.KeyValuesWriter;
	import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
	import org.apache.tez.runtime.library.api.TezRuntimeConfiguration.ReportPartitionStats;
	import org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler;
	import org.apache.tez.runtime.library.common.sort.impl.dflt.DefaultSorter;
	import org.apache.tez.runtime.library.conf.OrderedPartitionedKVOutputConfig.SorterImpl;
	import org.apache.tez.runtime.library.partitioner.HashPartitioner;
	import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads;
	import org.junit.After;
	import org.junit.Assert;
	import org.junit.Before;
	import org.junit.Test;
	import org.junit.runner.RunWith;
	import org.junit.runners.Parameterized;
	import org.mockito.invocation.InvocationOnMock;
	import org.mockito.stubbing.Answer;

	import java.io.IOException;
	import java.math.BigInteger;
	import java.nio.ByteBuffer;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.LinkedList;
	import java.util.List;
	import java.util.Random;
	import java.util.concurrent.atomic.AtomicLong;

	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertFalse;
	import static org.junit.Assert.assertTrue;
	import static org.junit.Assert.fail;
	import static org.mockito.Matchers.any;
	import static org.mockito.Matchers.anyLong;
	import static org.mockito.Mockito.doAnswer;
	import static org.mockito.Mockito.doReturn;
	import static org.mockito.Mockito.mock;

	@SuppressWarnings({ "rawtypes", "unchecked" })
	@RunWith(Parameterized.class)
	public class TestOnFileSortedOutput {
	private static final Random rnd = new Random();
	private static final String UniqueID = "UUID";
	private static final String HOST = "localhost";
	private static final int PORT = 80;

	private Configuration conf;
	private FileSystem fs;
	private Path workingDir;
	//no of outputs
	private int partitions;
	//For sorter (pipelined / Default)
	private SorterImpl sorterImpl;
	private int sorterThreads;

	final AtomicLong outputSize = new AtomicLong();
	final AtomicLong numRecords = new AtomicLong();

	private KeyValuesWriter writer;
	private OrderedPartitionedKVOutput sortedOutput;
	private boolean sendEmptyPartitionViaEvent;
	//Partition index for which data should not be written to.
	private int emptyPartitionIdx;
	private ReportPartitionStats reportPartitionStats;

	/**
	* Constructor
	*
	* @param sendEmptyPartitionViaEvent
	* @param sorterImpl Which sorter impl ( pipeline/legacy )
	* @param sorterThreads number of threads needed for sorter (required only for pipelined sorter)
	* @param emptyPartitionIdx for which data should not be generated
	*/
	public TestOnFileSortedOutput(boolean sendEmptyPartitionViaEvent,
	SorterImpl sorterImpl, int sorterThreads, int emptyPartitionIdx,
	ReportPartitionStats reportPartitionStats) throws IOException {
	this.sendEmptyPartitionViaEvent = sendEmptyPartitionViaEvent;
	this.emptyPartitionIdx = emptyPartitionIdx;
	this.sorterImpl = sorterImpl;
	this.sorterThreads = sorterThreads;
	this.reportPartitionStats = reportPartitionStats;
	conf = new Configuration();

	workingDir = new Path(".", this.getClass().getName());
	String localDirs = workingDir.toString();
	conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, localDirs);
	fs = FileSystem.getLocal(conf);
	}

	@Before
	public void setup() throws Exception {
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, sorterImpl.name());
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SORTER_SORT_THREADS, sorterThreads);
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 5);

	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS, Text.class.getName());
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS, Text.class.getName());
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS,
	HashPartitioner.class.getName());

	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_EMPTY_PARTITION_INFO_VIA_EVENTS_ENABLED,
	sendEmptyPartitionViaEvent);
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS,
	reportPartitionStats.getType());
	outputSize.set(0);
	numRecords.set(0);
	fs.mkdirs(workingDir);
	this.partitions = Math.max(1, rnd.nextInt(10));
	}

	@After
	public void cleanup() throws IOException {
	fs.delete(workingDir, true);
	}

	@SuppressWarnings("deprecation")
	@Parameterized.Parameters(name = "test[{0}, {1}, {2}, {3}, {4}]")
	public static Collection<Object[]> getParameters() {
	Collection<Object[]> parameters = new ArrayList<Object[]>();
	//empty_partition_via_events_enabled, noOfSortThreads,
	// partitionToBeEmpty, reportPartitionStats
	parameters.add(new Object[] { false, SorterImpl.LEGACY, 1, -1,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { false, SorterImpl.LEGACY, 1, 0,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, -1,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, 0,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.LEGACY, 1, 0,
	ReportPartitionStats.PRECISE });

	//Pipelined sorter
	parameters.add(new Object[] { false, SorterImpl.PIPELINED, 2, -1,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { false, SorterImpl.PIPELINED, 2, 0,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, -1,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, 0,
	ReportPartitionStats.ENABLED });
	parameters.add(new Object[] { true, SorterImpl.PIPELINED, 2, 0,
	ReportPartitionStats.PRECISE });
	return parameters;
	}

	private void startSortedOutput(int partitions) throws Exception {
	OutputContext context = createTezOutputContext();
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 4);
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
	sortedOutput.initialize();
	sortedOutput.start();
	writer = sortedOutput.getWriter();
	}

	private void _testPipelinedShuffle(String sorterName) throws Exception {
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);

	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, sorterName);
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);
	OutputContext context = createTezOutputContext();
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

	sortedOutput.initialize();
	sortedOutput.start();

	assertFalse(sortedOutput.finalMergeEnabled);
	assertTrue(sortedOutput.pipelinedShuffle);
	}

	@Test (timeout = 5000)
	public void testPipelinedShuffle() throws Exception {
	_testPipelinedShuffle(SorterImpl.PIPELINED.name());
	}

	@Test (timeout = 5000)
	public void testPipelinedShuffleWithFinalMerge() throws Exception {
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, SorterImpl.PIPELINED.name());

	//wrong setting for final merge enable in output
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, true);
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);
	OutputContext context = createTezOutputContext();
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

	sortedOutput.initialize();
	sortedOutput.start();
	assertFalse(sortedOutput.finalMergeEnabled); //should be disabled as pipelining is on
	assertTrue(sortedOutput.pipelinedShuffle);
	}

	@Test
	public void testPipelinedSettingsWithDefaultSorter() throws Exception {
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);
	//negative. with default sorter
	conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, SorterImpl.LEGACY.name());
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
	conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);

	OutputContext context = createTezOutputContext();
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

	sortedOutput.initialize();
	try {
	sortedOutput.start();
	fail("Should have thrown illegal argument exception as pipelining is enabled with "
	+ "DefaultSorter");
	} catch(IllegalArgumentException ie) {
	assertTrue(ie.getMessage().contains("works with PipelinedSorter"));
	}

	}

	@Test (timeout = 5000)
	public void testSortBufferSize() throws Exception{
	OutputContext context = createTezOutputContext();
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 2048);
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
	try {
	//Memory limit checks are done in sorter impls. For e.g, defaultsorter does not support > 2GB
	sortedOutput.initialize();
	DefaultSorter sorter = new DefaultSorter(context, conf, 100, 350010241024l);
	fail();
	} catch(IllegalArgumentException e) {
	assertTrue(e.getMessage().contains(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB));
	}
	conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 0);
	payLoad = TezUtils.createUserPayloadFromConf(conf);
	doReturn(payLoad).when(context).getUserPayload();
	sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
	try {
	sortedOutput.initialize();
	fail();
	} catch(IllegalArgumentException e) {
	assertTrue(e.getMessage().contains(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB));
	}
	}

	@Test(timeout = 5000)
	public void baseTest() throws Exception {
	startSortedOutput(partitions);

	//Write random set of keys
	long recordsWritten = numRecords.get();
	for (int i = 0; i < Math.max(1, rnd.nextInt(50)); i++) {
	Text key = new Text(new BigInteger(256, rnd).toString());
	LinkedList values = new LinkedList();
	for (int j = 0; j < Math.max(2, rnd.nextInt(10)); j++) {
	recordsWritten++;
	values.add(new Text(new BigInteger(256, rnd).toString()));
	}
	writer.write(key, values);
	}

	List<Event> eventList = sortedOutput.close();
	assertTrue(eventList != null && eventList.size() == 2);
	assertEquals(recordsWritten, numRecords.get());
	ShuffleUserPayloads.DataMovementEventPayloadProto
	payload = ShuffleUserPayloads.DataMovementEventPayloadProto
	.parseFrom(
	ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));

	ShuffleUserPayloads.VertexManagerEventPayloadProto
	vmPayload = ShuffleUserPayloads.VertexManagerEventPayloadProto
	.parseFrom(
	ByteString.copyFrom(((VertexManagerEvent) eventList.get(0)).getUserPayload()));

	if (reportPartitionStats.isPrecise()) {
	assertTrue(vmPayload.hasDetailedPartitionStats());
	} else {
	assertTrue(vmPayload.hasPartitionStats());
	}
	assertEquals(HOST, payload.getHost());
	assertEquals(PORT, payload.getPort());
	assertEquals(UniqueID, payload.getPathComponent());
	}

	@Test(timeout = 5000)
	public void testWithSomeEmptyPartition() throws Exception {
	//ensure atleast 2 partitions are available
	partitions = Math.max(2, partitions);
	startSortedOutput(partitions);

	//write random data
	for (int i = 0; i < 2 * partitions; i++) {
	Text key = new Text(new BigInteger(256, rnd).toString());
	Text value = new Text(new BigInteger(256, rnd).toString());
	//skip writing to certain partitions
	if (i % partitions != emptyPartitionIdx) {
	writer.write(key, value);
	}
	}

	List<Event> eventList = sortedOutput.close();
	assertTrue(eventList != null && eventList.size() == 2);

	ShuffleUserPayloads.DataMovementEventPayloadProto
	payload = ShuffleUserPayloads.DataMovementEventPayloadProto
	.parseFrom(ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));

	assertEquals(HOST, payload.getHost());
	assertEquals(PORT, payload.getPort());
	assertEquals(UniqueID, payload.getPathComponent());
	}

	@Test(timeout = 5000)
	public void testAllEmptyPartition() throws Exception {
	startSortedOutput(partitions);

	//Close output without writing any data to it.
	List<Event> eventList = sortedOutput.close();
	assertTrue(eventList != null && eventList.size() == 2);

	ShuffleUserPayloads.DataMovementEventPayloadProto
	payload = ShuffleUserPayloads.DataMovementEventPayloadProto
	.parseFrom(ByteString.copyFrom(((CompositeDataMovementEvent) eventList.get(1)).getUserPayload()));
	if (sendEmptyPartitionViaEvent) {
	assertEquals("", payload.getHost());
	assertEquals(0, payload.getPort());
	assertEquals("", payload.getPathComponent());
	} else {
	assertEquals(HOST, payload.getHost());
	assertEquals(PORT, payload.getPort());
	assertEquals(UniqueID, payload.getPathComponent());
	}
	}

	private OutputContext createTezOutputContext() throws IOException {
	String[] workingDirs = { workingDir.toString() };
	Configuration localConf = new Configuration(false);
	UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
	DataOutputBuffer serviceProviderMetaData = new DataOutputBuffer();
	serviceProviderMetaData.writeInt(PORT);

	TezCounters counters = new TezCounters();

	OutputStatisticsReporter reporter = mock(OutputStatisticsReporter.class);
	doAnswer(new Answer() {
	@Override public Object answer(InvocationOnMock invocation) throws Throwable {
	outputSize.set((Long) invocation.getArguments()[0]);
	return null;
	}
	}).when(reporter).reportDataSize(anyLong());
	doAnswer(new Answer() {
	@Override public Object answer(InvocationOnMock invocation) throws Throwable {
	numRecords.set((Long) invocation.getArguments()[0]);
	return null;
	}
	}).when(reporter).reportItemsProcessed(anyLong());


	OutputContext context = mock(OutputContext.class);
	doReturn(localConf).when(context).getContainerConfiguration();
	doReturn(counters).when(context).getCounters();
	doReturn(workingDirs).when(context).getWorkDirs();
	doReturn(payLoad).when(context).getUserPayload();
	doReturn(5 * 1024 * 1024l).when(context).getTotalMemoryAvailableToTask();
	doReturn(UniqueID).when(context).getUniqueIdentifier();
	doReturn("v0").when(context).getTaskVertexName();
	doReturn("v1").when(context).getDestinationVertexName();
	doReturn(ByteBuffer.wrap(serviceProviderMetaData.getData())).when(context)
	.getServiceProviderMetaData
	(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID,
	TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
	doAnswer(new Answer() {
	@Override public Object answer(InvocationOnMock invocation) throws Throwable {
	long requestedSize = (Long) invocation.getArguments()[0];
	MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation
	.getArguments()[1];
	callback.memoryAssigned(requestedSize);
	return null;
	}
	}).when(context).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
	ExecutionContext ExecutionContext = mock(ExecutionContext.class);
	doReturn(HOST).when(ExecutionContext).getHostName();
	doReturn(reporter).when(context).getStatisticsReporter();
	doReturn(ExecutionContext).when(context).getExecutionContext();
	return context;
	}

	@Test(timeout=5000)
	public void testInvalidSorter() throws Exception {
	try {
	_testPipelinedShuffle("Foo");
	Assert.fail("Expected start to fail due to invalid sorter");
	} catch (IllegalArgumentException e) {
	// Expected
	}
	}

	@Test(timeout=5000)
	public void testLowerCaseNamedSorter() throws Exception {
	_testPipelinedShuffle("Pipelined");
	}


	}