| /** |
| * 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.tez.runtime.library.common.writers; |
| |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertTrue; |
| import static org.mockito.Mockito.any; |
| import static org.mockito.Mockito.anyList; |
| import static org.mockito.Mockito.atLeast; |
| import static org.mockito.Mockito.atMost; |
| import static org.mockito.Mockito.doAnswer; |
| import static org.mockito.Mockito.doReturn; |
| import static org.mockito.Mockito.eq; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.never; |
| import static org.mockito.Mockito.times; |
| import static org.mockito.Mockito.verify; |
| |
| import java.io.ByteArrayInputStream; |
| import java.io.DataInputStream; |
| import java.io.InputStream; |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.util.Arrays; |
| import java.util.BitSet; |
| import java.util.Collection; |
| import java.util.HashMap; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Random; |
| import java.util.Set; |
| import java.util.UUID; |
| import java.util.regex.Matcher; |
| import java.util.regex.Pattern; |
| |
| import com.google.protobuf.ByteString; |
| import org.apache.hadoop.fs.CommonConfigurationKeys; |
| import org.apache.hadoop.fs.permission.FsAction; |
| import org.apache.tez.dag.api.TezConfiguration; |
| import org.apache.tez.runtime.api.events.VertexManagerEvent; |
| import org.apache.tez.runtime.library.common.Constants; |
| import org.apache.tez.runtime.library.common.writers.UnorderedPartitionedKVWriter.SpillInfo; |
| import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads; |
| import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads.VertexManagerEventPayloadProto; |
| import org.apache.tez.runtime.library.utils.DATA_RANGE_IN_MB; |
| import org.mockito.invocation.InvocationOnMock; |
| import org.roaringbitmap.RoaringBitmap; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| import org.apache.hadoop.conf.Configurable; |
| import org.apache.hadoop.conf.Configuration; |
| import org.apache.hadoop.fs.FSDataInputStream; |
| import org.apache.hadoop.fs.FileSystem; |
| import org.apache.hadoop.fs.Path; |
| import org.apache.hadoop.io.DataInputBuffer; |
| import org.apache.hadoop.io.IntWritable; |
| import org.apache.hadoop.io.LongWritable; |
| import org.apache.hadoop.io.Text; |
| import org.apache.hadoop.io.Writable; |
| import org.apache.hadoop.io.compress.CompressionCodec; |
| import org.apache.hadoop.io.compress.DefaultCodec; |
| import org.apache.hadoop.yarn.api.records.ApplicationId; |
| import org.apache.tez.common.TezCommonUtils; |
| import org.apache.tez.common.TezRuntimeFrameworkConfigs; |
| import org.apache.tez.common.TezUtilsInternal; |
| import org.apache.tez.common.counters.TaskCounter; |
| import org.apache.tez.common.counters.TezCounter; |
| import org.apache.tez.common.counters.TezCounters; |
| import org.apache.tez.runtime.api.Event; |
| import org.apache.tez.runtime.api.OutputContext; |
| import org.apache.tez.runtime.api.events.CompositeDataMovementEvent; |
| import org.apache.tez.runtime.library.api.Partitioner; |
| import org.apache.tez.runtime.library.api.TezRuntimeConfiguration; |
| import org.apache.tez.runtime.library.api.TezRuntimeConfiguration.ReportPartitionStats; |
| import org.apache.tez.runtime.library.common.sort.impl.IFile; |
| import org.apache.tez.runtime.library.common.sort.impl.TezIndexRecord; |
| import org.apache.tez.runtime.library.common.sort.impl.TezSpillRecord; |
| import org.apache.tez.runtime.library.common.task.local.output.TezTaskOutput; |
| import org.apache.tez.runtime.library.common.task.local.output.TezTaskOutputFiles; |
| import org.apache.tez.runtime.library.partitioner.HashPartitioner; |
| import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads.DataMovementEventPayloadProto; |
| import org.junit.After; |
| import org.junit.Before; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.Parameterized; |
| |
| import com.google.common.collect.LinkedListMultimap; |
| import com.google.common.collect.Multimap; |
| import com.google.common.collect.Sets; |
| import org.mockito.ArgumentCaptor; |
| import org.mockito.stubbing.Answer; |
| |
| @RunWith(value = Parameterized.class) |
| public class TestUnorderedPartitionedKVWriter { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(TestUnorderedPartitionedKVWriter.class); |
| |
| private static final String HOST_STRING = "localhost"; |
| private static final int SHUFFLE_PORT = 4000; |
| |
| private static String testTmpDir = System.getProperty("test.build.data", "/tmp"); |
| private static final Path TEST_ROOT_DIR = new Path(testTmpDir, |
| TestUnorderedPartitionedKVWriter.class.getSimpleName()); |
| private static FileSystem localFs; |
| |
| private boolean shouldCompress; |
| private ReportPartitionStats reportPartitionStats; |
| private Configuration defaultConf = new Configuration(); |
| |
| public TestUnorderedPartitionedKVWriter(boolean shouldCompress, |
| ReportPartitionStats reportPartitionStats) { |
| this.shouldCompress = shouldCompress; |
| this.reportPartitionStats = reportPartitionStats; |
| } |
| |
| @SuppressWarnings("deprecation") |
| @Parameterized.Parameters(name = "test[{0}, {1}]") |
| public static Collection<Object[]> data() { |
| Object[][] data = new Object[][] { |
| { false, ReportPartitionStats.DISABLED }, |
| { false, ReportPartitionStats.ENABLED }, |
| { false, ReportPartitionStats.NONE }, |
| { false, ReportPartitionStats.MEMORY_OPTIMIZED }, |
| { false, ReportPartitionStats.PRECISE }, |
| { true, ReportPartitionStats.DISABLED }, |
| { true, ReportPartitionStats.ENABLED }, |
| { true, ReportPartitionStats.NONE }, |
| { true, ReportPartitionStats.MEMORY_OPTIMIZED }, |
| { true, ReportPartitionStats.PRECISE }}; |
| return Arrays.asList(data); |
| } |
| |
| @Before |
| public void setup() throws IOException { |
| LOG.info("Setup. Using test dir: " + TEST_ROOT_DIR); |
| localFs = FileSystem.getLocal(new Configuration()); |
| localFs.delete(TEST_ROOT_DIR, true); |
| localFs.mkdirs(TEST_ROOT_DIR); |
| } |
| |
| @After |
| public void cleanup() throws IOException { |
| LOG.info("CleanUp"); |
| localFs.delete(TEST_ROOT_DIR, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testBufferSizing() throws IOException { |
| ApplicationId appId = ApplicationId.newInstance(10000000, 1); |
| TezCounters counters = new TezCounters(); |
| String uniqueId = UUID.randomUUID().toString(); |
| String auxiliaryService = defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, |
| TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT); |
| OutputContext outputContext = createMockOutputContext(counters, appId, uniqueId, auxiliaryService); |
| |
| final int maxSingleBufferSizeBytes = 2047; |
| final long sizePerBuffer = maxSingleBufferSizeBytes - 64 - maxSingleBufferSizeBytes % 4; |
| Configuration conf = createConfiguration(outputContext, IntWritable.class, LongWritable.class, |
| false, maxSingleBufferSizeBytes); |
| |
| int numOutputs = 10; |
| |
| UnorderedPartitionedKVWriter kvWriter = null; |
| |
| // Not enough memory so divide into 2 buffers. |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, 2048); |
| assertEquals(2, kvWriter.numBuffers); |
| assertEquals(1024, kvWriter.sizePerBuffer); |
| assertEquals(1024, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| |
| // allocate exact |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| maxSingleBufferSizeBytes * 3); |
| assertEquals(3, kvWriter.numBuffers); |
| assertEquals(sizePerBuffer, kvWriter.sizePerBuffer); |
| assertEquals(sizePerBuffer, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| |
| // under allocate |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| maxSingleBufferSizeBytes * 2 + maxSingleBufferSizeBytes / 2); |
| assertEquals(2, kvWriter.numBuffers); |
| assertEquals(sizePerBuffer, kvWriter.sizePerBuffer); |
| assertEquals(sizePerBuffer, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| |
| // over allocate |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| maxSingleBufferSizeBytes * 2 + maxSingleBufferSizeBytes / 2 + 1); |
| assertEquals(3, kvWriter.numBuffers); |
| assertEquals(sizePerBuffer, kvWriter.sizePerBuffer); |
| assertEquals(maxSingleBufferSizeBytes / 2 + 1, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| |
| // spill limit 1. |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| 4 * maxSingleBufferSizeBytes + 1); |
| assertEquals(4, kvWriter.numBuffers); |
| assertEquals(sizePerBuffer, kvWriter.sizePerBuffer); |
| assertEquals(sizePerBuffer, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| |
| // spill limit 2. |
| conf.setInt( |
| TezRuntimeConfiguration.TEZ_RUNTIME_UNORDERED_PARTITIONED_KVWRITER_BUFFER_MERGE_PERCENT, |
| 50); |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| 4 * maxSingleBufferSizeBytes + 1); |
| assertEquals(4, kvWriter.numBuffers); |
| assertEquals(sizePerBuffer, kvWriter.sizePerBuffer); |
| assertEquals(sizePerBuffer, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(2, kvWriter.spillLimit); |
| |
| // Available memory is less than buffer size. |
| conf.unset(TezRuntimeConfiguration.TEZ_RUNTIME_UNORDERED_OUTPUT_MAX_PER_BUFFER_SIZE_BYTES); |
| kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, conf, numOutputs, |
| 2048); |
| assertEquals(2, kvWriter.numBuffers); |
| assertEquals(1024, kvWriter.sizePerBuffer); |
| assertEquals(1024, kvWriter.lastBufferSize); |
| assertEquals(1, kvWriter.numInitializedBuffers); |
| assertEquals(1, kvWriter.spillLimit); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoSpill() throws IOException, InterruptedException { |
| baseTest(10, 10, null, shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSingleSpill() throws IOException, InterruptedException { |
| baseTest(50, 10, null, shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testMultipleSpills() throws IOException, InterruptedException { |
| baseTest(200, 10, null, shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testMultipleSpillsWithSmallBuffer() throws IOException, InterruptedException { |
| // numBuffers is much higher than available threads. |
| baseTest(200, 10, null, shouldCompress, 512, 0, 9600, false); |
| } |
| |
| @Test(timeout = 10000) |
| public void testMergeBuffersAndSpill() throws IOException, InterruptedException { |
| baseTest(200, 10, null, shouldCompress, 2048, 10); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords() throws IOException, InterruptedException { |
| baseTest(0, 10, null, shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords_SinglePartition() throws IOException, InterruptedException { |
| // skipBuffers |
| baseTest(0, 1, null, shouldCompress, -1, 0, 2048, false); |
| // Check with data via events |
| baseTest(0, 1, null, shouldCompress, -1, 0, 2048, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSkippedPartitions() throws IOException, InterruptedException { |
| baseTest(200, 10, Sets.newHashSet(2, 5), shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoSpill_SinglePartition() throws IOException, InterruptedException { |
| baseTest(10, 1, null, shouldCompress, -1, 0); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSpill_SinglePartition() throws IOException, InterruptedException { |
| baseTest(1000, 1, null, shouldCompress, -1, 0, 2048, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testRandomText() throws IOException, InterruptedException { |
| textTest(100, 10, 2048, 0, 0, 0, false, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargeKeys() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 10, 0, 0, false, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargevalues() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 0, 10, 0, false, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargeKvPairs() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 0, 0, 10, false, true); |
| } |
| |
| @Test(timeout = 10000) |
| public void testTextMixedRecords() throws IOException, InterruptedException { |
| textTest(100, 10, 2048, 10, 10, 10, false, true); |
| } |
| |
| @Test(timeout = 10000000) |
| public void testRandomTextWithoutFinalMerge() throws IOException, InterruptedException { |
| textTest(100, 10, 2048, 0, 0, 0, false, false); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargeKeysWithoutFinalMerge() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 10, 0, 0, false, false); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargevaluesWithoutFinalMerge() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 0, 10, 0, false, false); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargeKvPairsWithoutFinalMerge() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 0, 0, 10, false, false); |
| } |
| |
| @Test(timeout = 10000) |
| public void testTextMixedRecordsWithoutFinalMerge() throws IOException, InterruptedException { |
| textTest(100, 10, 2048, 10, 10, 10, false, false); |
| } |
| |
| public void textTest(int numRegularRecords, int numPartitions, long availableMemory, |
| int numLargeKeys, int numLargevalues, int numLargeKvPairs, |
| boolean pipeliningEnabled, boolean isFinalMergeEnabled) throws IOException, |
| InterruptedException { |
| Partitioner partitioner = new HashPartitioner(); |
| ApplicationId appId = ApplicationId.newInstance(10000000, 1); |
| TezCounters counters = new TezCounters(); |
| String uniqueId = UUID.randomUUID().toString(); |
| int dagId = 1; |
| String auxiliaryService = defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, |
| TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT); |
| OutputContext outputContext = createMockOutputContext(counters, appId, uniqueId, auxiliaryService); |
| Random random = new Random(); |
| |
| Configuration conf = createConfiguration(outputContext, Text.class, Text.class, shouldCompress, |
| -1, HashPartitioner.class); |
| conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, pipeliningEnabled); |
| conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, isFinalMergeEnabled); |
| |
| CompressionCodec codec = null; |
| if (shouldCompress) { |
| codec = new DefaultCodec(); |
| ((Configurable) codec).setConf(conf); |
| } |
| |
| int numRecordsWritten = 0; |
| |
| Map<Integer, Multimap<String, String>> expectedValues = new HashMap<Integer, Multimap<String, String>>(); |
| for (int i = 0; i < numPartitions; i++) { |
| expectedValues.put(i, LinkedListMultimap.<String, String> create()); |
| } |
| |
| UnorderedPartitionedKVWriter kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, |
| conf, numPartitions, availableMemory); |
| |
| int sizePerBuffer = kvWriter.sizePerBuffer; |
| |
| BitSet partitionsWithData = new BitSet(numPartitions); |
| Text keyText = new Text(); |
| Text valText = new Text(); |
| for (int i = 0; i < numRegularRecords; i++) { |
| String key = createRandomString(Math.abs(random.nextInt(10))); |
| String val = createRandomString(Math.abs(random.nextInt(20))); |
| keyText.set(key); |
| valText.set(val); |
| int partition = partitioner.getPartition(keyText, valText, numPartitions); |
| partitionsWithData.set(partition); |
| expectedValues.get(partition).put(key, val); |
| kvWriter.write(keyText, valText); |
| numRecordsWritten++; |
| } |
| |
| // Write Large key records |
| for (int i = 0; i < numLargeKeys; i++) { |
| String key = createRandomString(sizePerBuffer + Math.abs(random.nextInt(100))); |
| String val = createRandomString(Math.abs(random.nextInt(20))); |
| keyText.set(key); |
| valText.set(val); |
| int partition = partitioner.getPartition(keyText, valText, numPartitions); |
| partitionsWithData.set(partition); |
| expectedValues.get(partition).put(key, val); |
| kvWriter.write(keyText, valText); |
| numRecordsWritten++; |
| } |
| if (pipeliningEnabled) { |
| verify(outputContext, times(numLargeKeys)).sendEvents(anyList()); |
| } |
| |
| // Write Large val records |
| for (int i = 0; i < numLargevalues; i++) { |
| String key = createRandomString(Math.abs(random.nextInt(10))); |
| String val = createRandomString(sizePerBuffer + Math.abs(random.nextInt(100))); |
| keyText.set(key); |
| valText.set(val); |
| int partition = partitioner.getPartition(keyText, valText, numPartitions); |
| partitionsWithData.set(partition); |
| expectedValues.get(partition).put(key, val); |
| kvWriter.write(keyText, valText); |
| numRecordsWritten++; |
| } |
| if (pipeliningEnabled) { |
| verify(outputContext, times(numLargevalues + numLargeKeys)).sendEvents(anyList()); |
| } |
| |
| // Write records where key + val are large (but both can fit in the buffer individually) |
| for (int i = 0; i < numLargeKvPairs; i++) { |
| String key = createRandomString(sizePerBuffer / 2 + Math.abs(random.nextInt(100))); |
| String val = createRandomString(sizePerBuffer / 2 + Math.abs(random.nextInt(100))); |
| keyText.set(key); |
| valText.set(val); |
| int partition = partitioner.getPartition(keyText, valText, numPartitions); |
| partitionsWithData.set(partition); |
| expectedValues.get(partition).put(key, val); |
| kvWriter.write(keyText, valText); |
| numRecordsWritten++; |
| } |
| if (pipeliningEnabled) { |
| verify(outputContext, times(numLargevalues + numLargeKeys + numLargeKvPairs)) |
| .sendEvents(anyList()); |
| } |
| |
| List<Event> events = kvWriter.close(); |
| verify(outputContext, never()).reportFailure(any(), any(), any()); |
| |
| if (!pipeliningEnabled) { |
| VertexManagerEvent vmEvent = null; |
| for (Event event : events) { |
| if (event instanceof VertexManagerEvent) { |
| assertNull(vmEvent); |
| vmEvent = (VertexManagerEvent) event; |
| } |
| } |
| VertexManagerEventPayloadProto vmEventPayload = |
| VertexManagerEventPayloadProto.parseFrom( |
| ByteString.copyFrom(vmEvent.getUserPayload().asReadOnlyBuffer())); |
| assertEquals(numRecordsWritten, vmEventPayload.getNumRecord()); |
| } |
| |
| TezCounter outputLargeRecordsCounter = counters.findCounter(TaskCounter.OUTPUT_LARGE_RECORDS); |
| assertEquals(numLargeKeys + numLargevalues + numLargeKvPairs, |
| outputLargeRecordsCounter.getValue()); |
| |
| if (pipeliningEnabled || !isFinalMergeEnabled) { |
| // verify spill data files and index file exist |
| for (int i = 0; i < kvWriter.numSpills.get(); i++) { |
| assertTrue(localFs.exists(kvWriter.outputFileHandler.getSpillFileForWrite(i, 0))); |
| assertTrue(localFs.exists(kvWriter.outputFileHandler.getSpillIndexFileForWrite(i, 0))); |
| } |
| return; |
| } |
| |
| // Validate the events |
| assertEquals(2, events.size()); |
| |
| assertTrue(events.get(0) instanceof VertexManagerEvent); |
| VertexManagerEvent vme = (VertexManagerEvent) events.get(0); |
| verifyPartitionStats(vme, partitionsWithData); |
| |
| assertTrue(events.get(1) instanceof CompositeDataMovementEvent); |
| CompositeDataMovementEvent cdme = (CompositeDataMovementEvent) events.get(1); |
| assertEquals(0, cdme.getSourceIndexStart()); |
| assertEquals(numPartitions, cdme.getCount()); |
| DataMovementEventPayloadProto eventProto = DataMovementEventPayloadProto.parseFrom( |
| ByteString.copyFrom(cdme.getUserPayload())); |
| assertFalse(eventProto.hasData()); |
| BitSet emptyPartitionBits = null; |
| if (partitionsWithData.cardinality() != numPartitions) { |
| assertTrue(eventProto.hasEmptyPartitions()); |
| byte[] emptyPartitions = TezCommonUtils.decompressByteStringToByteArray( |
| eventProto.getEmptyPartitions()); |
| emptyPartitionBits = TezUtilsInternal.fromByteArray(emptyPartitions); |
| assertEquals(numPartitions - partitionsWithData.cardinality(), |
| emptyPartitionBits.cardinality()); |
| } else { |
| assertFalse(eventProto.hasEmptyPartitions()); |
| emptyPartitionBits = new BitSet(numPartitions); |
| } |
| assertEquals(HOST_STRING, eventProto.getHost()); |
| assertEquals(SHUFFLE_PORT, eventProto.getPort()); |
| assertEquals(uniqueId, eventProto.getPathComponent()); |
| |
| // Verify the data |
| // Verify the actual data |
| TezTaskOutput taskOutput = new TezTaskOutputFiles(conf, uniqueId, dagId); |
| Path outputFilePath = kvWriter.finalOutPath; |
| Path spillFilePath = kvWriter.finalIndexPath; |
| if (numRecordsWritten > 0) { |
| assertTrue(localFs.exists(outputFilePath)); |
| assertTrue(localFs.exists(spillFilePath)); |
| checkPermissions(outputFilePath, spillFilePath); |
| } else { |
| return; |
| } |
| |
| // Special case for 0 records. |
| TezSpillRecord spillRecord = new TezSpillRecord(spillFilePath, conf); |
| DataInputBuffer keyBuffer = new DataInputBuffer(); |
| DataInputBuffer valBuffer = new DataInputBuffer(); |
| Text keyDeser = new Text(); |
| Text valDeser = new Text(); |
| for (int i = 0; i < numPartitions; i++) { |
| if (emptyPartitionBits.get(i)) { |
| continue; |
| } |
| TezIndexRecord indexRecord = spillRecord.getIndex(i); |
| FSDataInputStream inStream = FileSystem.getLocal(conf).open(outputFilePath); |
| inStream.seek(indexRecord.getStartOffset()); |
| IFile.Reader reader = new IFile.Reader(inStream, indexRecord.getPartLength(), codec, null, |
| null, false, 0, -1); |
| while (reader.nextRawKey(keyBuffer)) { |
| reader.nextRawValue(valBuffer); |
| keyDeser.readFields(keyBuffer); |
| valDeser.readFields(valBuffer); |
| int partition = partitioner.getPartition(keyDeser, valDeser, numPartitions); |
| assertTrue(expectedValues.get(partition).remove(keyDeser.toString(), valDeser.toString())); |
| } |
| inStream.close(); |
| } |
| for (int i = 0; i < numPartitions; i++) { |
| assertEquals(0, expectedValues.get(i).size()); |
| expectedValues.remove(i); |
| } |
| assertEquals(0, expectedValues.size()); |
| } |
| |
| private int[] getPartitionStats(VertexManagerEvent vme) throws IOException { |
| RoaringBitmap partitionStats = new RoaringBitmap(); |
| VertexManagerEventPayloadProto |
| payload = VertexManagerEventPayloadProto |
| .parseFrom(ByteString.copyFrom(vme.getUserPayload())); |
| if (!reportPartitionStats.isEnabled()) { |
| assertFalse(payload.hasPartitionStats()); |
| assertFalse(payload.hasDetailedPartitionStats()); |
| return null; |
| } |
| if (reportPartitionStats.isPrecise()) { |
| assertTrue(payload.hasDetailedPartitionStats()); |
| List<Integer> sizeInMBList = |
| payload.getDetailedPartitionStats().getSizeInMbList(); |
| int[] stats = new int[sizeInMBList.size()]; |
| for (int i=0; i<sizeInMBList.size(); i++) { |
| stats[i] += sizeInMBList.get(i); |
| } |
| return stats; |
| } else { |
| assertTrue(payload.hasPartitionStats()); |
| ByteString compressedPartitionStats = payload.getPartitionStats(); |
| byte[] rawData = TezCommonUtils.decompressByteStringToByteArray( |
| compressedPartitionStats); |
| ByteArrayInputStream bin = new ByteArrayInputStream(rawData); |
| partitionStats.deserialize(new DataInputStream(bin)); |
| int[] stats = new int[partitionStats.getCardinality()]; |
| Iterator<Integer> it = partitionStats.iterator(); |
| final DATA_RANGE_IN_MB[] RANGES = DATA_RANGE_IN_MB.values(); |
| final int RANGE_LEN = RANGES.length; |
| while (it.hasNext()) { |
| int pos = it.next(); |
| int index = ((pos) / RANGE_LEN); |
| int rangeIndex = ((pos) % RANGE_LEN); |
| if (RANGES[rangeIndex].getSizeInMB() > 0) { |
| stats[index] += RANGES[rangeIndex].getSizeInMB(); |
| } |
| } |
| return stats; |
| } |
| } |
| |
| private void verifyPartitionStats(VertexManagerEvent vme, |
| BitSet expectedPartitionsWithData) throws IOException { |
| int[] stats = getPartitionStats(vme); |
| if (stats == null) { |
| return; |
| } |
| for (int i = 0; i < stats.length; i++) { |
| // The stats should be greater than zero if and only if |
| // the partition has data |
| assertTrue(expectedPartitionsWithData.get(i) == (stats[i] > 0)); |
| } |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoSpill_WithPipelinedShuffle() throws IOException, InterruptedException { |
| baseTestWithPipelinedTransfer(10, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSingleSpill_WithPipelinedShuffle() throws IOException, InterruptedException { |
| baseTestWithPipelinedTransfer(50, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testMultipleSpills_WithPipelinedShuffle() throws IOException, InterruptedException { |
| baseTestWithPipelinedTransfer(200, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords_WithPipelinedShuffle() throws IOException, InterruptedException { |
| baseTestWithPipelinedTransfer(0, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords_SinglePartition_WithPipelinedShuffle() throws IOException, InterruptedException { |
| // skipBuffers |
| baseTestWithPipelinedTransfer(0, 1, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSkippedPartitions_WithPipelinedShuffle() throws IOException, InterruptedException { |
| baseTestWithPipelinedTransfer(200, 10, Sets.newHashSet(2, 5), shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testLargeKvPairs_WithPipelinedShuffle() throws IOException, InterruptedException { |
| textTest(0, 10, 2048, 10, 20, 50, true, false); |
| } |
| |
| |
| @SuppressWarnings("unchecked") |
| private void baseTestWithPipelinedTransfer(int numRecords, int numPartitions, Set<Integer> |
| skippedPartitions, boolean shouldCompress) throws IOException, InterruptedException { |
| |
| PartitionerForTest partitioner = new PartitionerForTest(); |
| ApplicationId appId = ApplicationId.newInstance(10000000, 1); |
| TezCounters counters = new TezCounters(); |
| String uniqueId = UUID.randomUUID().toString(); |
| int dagId = 1; |
| String auxiliaryService = defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, |
| TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT); |
| OutputContext outputContext = createMockOutputContext(counters, appId, uniqueId, auxiliaryService); |
| |
| Configuration conf = createConfiguration(outputContext, IntWritable.class, LongWritable.class, |
| shouldCompress, -1); |
| conf.setBoolean(TezRuntimeConfiguration |
| .TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false); |
| conf.setBoolean(TezRuntimeConfiguration |
| .TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true); |
| |
| CompressionCodec codec = null; |
| if (shouldCompress) { |
| codec = new DefaultCodec(); |
| ((Configurable) codec).setConf(conf); |
| } |
| |
| int numOutputs = numPartitions; |
| long availableMemory = 2048; |
| int numRecordsWritten = 0; |
| |
| UnorderedPartitionedKVWriter kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, |
| conf, numOutputs, availableMemory); |
| |
| int sizePerBuffer = kvWriter.sizePerBuffer; |
| int sizePerRecord = 4 + 8; // IntW + LongW |
| int sizePerRecordWithOverhead = sizePerRecord + 12; // Record + META_OVERHEAD |
| |
| BitSet partitionsWithData = new BitSet(numPartitions); |
| IntWritable intWritable = new IntWritable(); |
| LongWritable longWritable = new LongWritable(); |
| for (int i = 0; i < numRecords; i++) { |
| intWritable.set(i); |
| longWritable.set(i); |
| int partition = partitioner.getPartition(intWritable, longWritable, numOutputs); |
| if (skippedPartitions != null && skippedPartitions.contains(partition)) { |
| continue; |
| } |
| partitionsWithData.set(partition); |
| kvWriter.write(intWritable, longWritable); |
| numRecordsWritten++; |
| } |
| |
| int recordsPerBuffer = sizePerBuffer / sizePerRecordWithOverhead; |
| int numExpectedSpills = numRecordsWritten / recordsPerBuffer; |
| |
| ArgumentCaptor<List> eventCaptor = ArgumentCaptor.forClass(List.class); |
| List<Event> lastEvents = kvWriter.close(); |
| |
| if (numPartitions == 1) { |
| assertEquals(false, kvWriter.skipBuffers); |
| } |
| |
| //no events are sent to kvWriter upon close with pipelining |
| assertTrue(lastEvents.size() == 0); |
| verify(outputContext, atLeast(numExpectedSpills)).sendEvents(eventCaptor.capture()); |
| int numOfCapturedEvents = eventCaptor.getAllValues().size(); |
| lastEvents = eventCaptor.getAllValues().get(numOfCapturedEvents - 1); |
| VertexManagerEvent VMEvent = (VertexManagerEvent)lastEvents.get(0); |
| |
| for (int i=0; i<numOfCapturedEvents; i++) { |
| List<Event> events = eventCaptor.getAllValues().get(i); |
| if (i < numOfCapturedEvents - 1) { |
| assertTrue(events.size() == 1); |
| assertTrue(events.get(0) instanceof CompositeDataMovementEvent); |
| } else { |
| assertTrue(events.size() == 2); |
| assertTrue(events.get(0) instanceof VertexManagerEvent); |
| assertTrue(events.get(1) instanceof CompositeDataMovementEvent); |
| } |
| } |
| verifyPartitionStats(VMEvent, partitionsWithData); |
| |
| verify(outputContext, never()).reportFailure(any(), |
| any(), any()); |
| |
| assertNull(kvWriter.currentBuffer); |
| assertEquals(0, kvWriter.availableBuffers.size()); |
| |
| // Verify the counters |
| TezCounter outputRecordBytesCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES); |
| TezCounter outputRecordsCounter = |
| counters.findCounter(TaskCounter.OUTPUT_RECORDS); |
| TezCounter outputBytesWithOverheadCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES_WITH_OVERHEAD); |
| TezCounter fileOutputBytesCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES_PHYSICAL); |
| TezCounter spilledRecordsCounter = |
| counters.findCounter(TaskCounter.SPILLED_RECORDS); |
| TezCounter additionalSpillBytesWritternCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_WRITTEN); |
| TezCounter additionalSpillBytesReadCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_READ); |
| TezCounter numAdditionalSpillsCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILL_COUNT); |
| assertEquals(numRecordsWritten * sizePerRecord, |
| outputRecordBytesCounter.getValue()); |
| assertEquals(numRecordsWritten, outputRecordsCounter.getValue()); |
| assertEquals(numRecordsWritten * sizePerRecordWithOverhead, |
| outputBytesWithOverheadCounter.getValue()); |
| long fileOutputBytes = fileOutputBytesCounter.getValue(); |
| if (numRecordsWritten > 0) { |
| assertTrue(fileOutputBytes > 0); |
| if (!shouldCompress) { |
| assertTrue(fileOutputBytes > outputRecordBytesCounter.getValue()); |
| } |
| } else { |
| assertEquals(0, fileOutputBytes); |
| } |
| // due to multiple threads, buffers could be merged in chunks in scheduleSpill. |
| assertTrue(recordsPerBuffer * numExpectedSpills >= spilledRecordsCounter.getValue()); |
| long additionalSpillBytesWritten = |
| additionalSpillBytesWritternCounter.getValue(); |
| long additionalSpillBytesRead = additionalSpillBytesReadCounter.getValue(); |
| |
| //No additional spill bytes written when final merge is disabled. |
| assertEquals(additionalSpillBytesWritten, 0); |
| |
| //No additional spills when final merge is disabled. |
| assertTrue(additionalSpillBytesWritten == additionalSpillBytesRead); |
| |
| //No additional spills when final merge is disabled. |
| assertEquals(numAdditionalSpillsCounter.getValue(), 0); |
| |
| assertTrue(lastEvents.size() > 0); |
| //Get the last event |
| int index = lastEvents.size() - 1; |
| assertTrue(lastEvents.get(index) instanceof CompositeDataMovementEvent); |
| CompositeDataMovementEvent cdme = |
| (CompositeDataMovementEvent)lastEvents.get(index); |
| assertEquals(0, cdme.getSourceIndexStart()); |
| assertEquals(numOutputs, cdme.getCount()); |
| DataMovementEventPayloadProto eventProto = |
| DataMovementEventPayloadProto.parseFrom( |
| ByteString.copyFrom(cdme.getUserPayload())); |
| //Ensure that this is the last event |
| assertTrue(eventProto.getLastEvent()); |
| verifyEmptyPartitions(eventProto, numRecordsWritten, numPartitions, skippedPartitions); |
| |
| verify(outputContext, atLeast(1)).notifyProgress(); |
| |
| // Verify if all spill files are available. |
| TezTaskOutput taskOutput = new TezTaskOutputFiles(conf, uniqueId, dagId); |
| |
| if (numRecordsWritten > 0) { |
| int numSpills = kvWriter.numSpills.get(); |
| for (int i = 0; i < numSpills; i++) { |
| Path outputFile = taskOutput.getSpillFileForWrite(i, 10); |
| Path indexFile = taskOutput.getSpillIndexFileForWrite(i, 10); |
| assertTrue(localFs.exists(outputFile)); |
| assertTrue(localFs.exists(indexFile)); |
| checkPermissions(outputFile, indexFile); |
| } |
| } else { |
| return; |
| } |
| } |
| |
| private void checkPermissions(Path outputFile, Path indexFile) throws IOException { |
| assertEquals("Incorrect output permissions (user)", FsAction.READ_WRITE, |
| localFs.getFileStatus(outputFile).getPermission().getUserAction()); |
| assertEquals("Incorrect output permissions (group)", FsAction.READ, |
| localFs.getFileStatus(outputFile).getPermission().getGroupAction()); |
| assertEquals("Incorrect index permissions (user)", FsAction.READ_WRITE, |
| localFs.getFileStatus(indexFile).getPermission().getUserAction()); |
| assertEquals("Incorrect index permissions (group)", FsAction.READ, |
| localFs.getFileStatus(indexFile).getPermission().getGroupAction()); |
| } |
| |
| private void verifyEmptyPartitions(DataMovementEventPayloadProto eventProto, |
| int numRecordsWritten, int numPartitions, Set<Integer> skippedPartitions) |
| throws IOException { |
| if (eventProto.hasEmptyPartitions()) { |
| byte[] emptyPartitions = TezCommonUtils.decompressByteStringToByteArray( |
| eventProto.getEmptyPartitions()); |
| BitSet emptyPartitionBits = TezUtilsInternal.fromByteArray(emptyPartitions); |
| if (numRecordsWritten == 0) { |
| assertEquals(numPartitions, emptyPartitionBits.cardinality()); |
| } else { |
| if (skippedPartitions != null) { |
| for (Integer e : skippedPartitions) { |
| assertTrue(emptyPartitionBits.get(e)); |
| } |
| assertEquals(skippedPartitions.size(), emptyPartitionBits.cardinality()); |
| } |
| } |
| if (emptyPartitionBits.cardinality() != numPartitions) { |
| assertEquals(HOST_STRING, eventProto.getHost()); |
| assertEquals(SHUFFLE_PORT, eventProto.getPort()); |
| assertTrue(eventProto.hasPathComponent()); |
| } else { |
| assertFalse(eventProto.hasHost()); |
| assertFalse(eventProto.hasPort()); |
| assertFalse(eventProto.hasPathComponent()); |
| } |
| } else { |
| assertEquals(HOST_STRING, eventProto.getHost()); |
| assertEquals(SHUFFLE_PORT, eventProto.getPort()); |
| assertTrue(eventProto.hasPathComponent()); |
| } |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoSpill_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(10, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSingleSpill_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(50, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSinglePartition_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(0, 1, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testMultipleSpills_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(200, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(0, 10, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testNoRecords_SinglePartition_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(0, 1, null, shouldCompress); |
| } |
| |
| @Test(timeout = 10000) |
| public void testSkippedPartitions_WithFinalMergeDisabled() throws IOException, InterruptedException { |
| baseTestWithFinalMergeDisabled(200, 10, Sets.newHashSet(2, 5), shouldCompress); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void baseTestWithFinalMergeDisabled(int numRecords, int numPartitions, |
| Set<Integer> skippedPartitions, boolean shouldCompress) throws IOException, InterruptedException { |
| |
| PartitionerForTest partitioner = new PartitionerForTest(); |
| ApplicationId appId = ApplicationId.newInstance(10000000, 1); |
| TezCounters counters = new TezCounters(); |
| String uniqueId = UUID.randomUUID().toString(); |
| int dagId = 1; |
| String auxiliaryService = defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, |
| TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT); |
| OutputContext outputContext = createMockOutputContext(counters, appId, uniqueId, auxiliaryService); |
| |
| Configuration conf = createConfiguration(outputContext, IntWritable.class, LongWritable.class, |
| shouldCompress, -1); |
| conf.setBoolean(TezRuntimeConfiguration |
| .TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false); |
| conf.setBoolean(TezRuntimeConfiguration |
| .TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, false); |
| |
| CompressionCodec codec = null; |
| if (shouldCompress) { |
| codec = new DefaultCodec(); |
| ((Configurable) codec).setConf(conf); |
| } |
| |
| int numOutputs = numPartitions; |
| long availableMemory = 2048; |
| int numRecordsWritten = 0; |
| |
| UnorderedPartitionedKVWriter kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, |
| conf, numOutputs, availableMemory); |
| |
| int sizePerBuffer = kvWriter.sizePerBuffer; |
| int sizePerRecord = 4 + 8; // IntW + LongW |
| int sizePerRecordWithOverhead = sizePerRecord + 12; // Record + META_OVERHEAD |
| |
| BitSet partitionsWithData = new BitSet(numPartitions); |
| IntWritable intWritable = new IntWritable(); |
| LongWritable longWritable = new LongWritable(); |
| for (int i = 0; i < numRecords; i++) { |
| intWritable.set(i); |
| longWritable.set(i); |
| int partition = partitioner.getPartition(intWritable, longWritable, numOutputs); |
| if (skippedPartitions != null && skippedPartitions.contains(partition)) { |
| continue; |
| } |
| partitionsWithData.set(partition); |
| kvWriter.write(intWritable, longWritable); |
| numRecordsWritten++; |
| } |
| |
| int recordsPerBuffer = sizePerBuffer / sizePerRecordWithOverhead; |
| int numExpectedSpills = numRecordsWritten / recordsPerBuffer; |
| |
| ArgumentCaptor<List> eventCaptor = ArgumentCaptor.forClass(List.class); |
| List<Event> lastEvents = kvWriter.close(); |
| |
| if (numPartitions == 1) { |
| assertEquals(true, kvWriter.skipBuffers); |
| } |
| |
| // max events sent are spills + one VM event. If there are no spills, atleast empty |
| // partitions would be sent out finally. |
| int spills = Math.max(1, kvWriter.numSpills.get()); |
| assertEquals((spills + 1), lastEvents.size()); //spills + VMEvent |
| verify(outputContext, atMost(0)).sendEvents(eventCaptor.capture()); |
| |
| for (int i=0; i<lastEvents.size(); i++) { |
| Event event =lastEvents.get(i); |
| if (event instanceof VertexManagerEvent) { |
| //when there are no records, empty IFile with 6 bytes would be created which would add up |
| // to stats. |
| if (numRecordsWritten > 0) { |
| verifyPartitionStats(((VertexManagerEvent) event), partitionsWithData); |
| } |
| } |
| } |
| |
| verify(outputContext, never()).reportFailure(any(), |
| any(), any()); |
| |
| assertNull(kvWriter.currentBuffer); |
| assertEquals(0, kvWriter.availableBuffers.size()); |
| |
| // Verify the counters |
| TezCounter outputRecordBytesCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES); |
| TezCounter outputRecordsCounter = |
| counters.findCounter(TaskCounter.OUTPUT_RECORDS); |
| TezCounter outputBytesWithOverheadCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES_WITH_OVERHEAD); |
| TezCounter fileOutputBytesCounter = |
| counters.findCounter(TaskCounter.OUTPUT_BYTES_PHYSICAL); |
| TezCounter spilledRecordsCounter = |
| counters.findCounter(TaskCounter.SPILLED_RECORDS); |
| TezCounter additionalSpillBytesWritternCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_WRITTEN); |
| TezCounter additionalSpillBytesReadCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_READ); |
| TezCounter numAdditionalSpillsCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILL_COUNT); |
| assertEquals(numRecordsWritten * sizePerRecord, |
| outputRecordBytesCounter.getValue()); |
| assertEquals(numRecordsWritten, outputRecordsCounter.getValue()); |
| if (outputRecordsCounter.getValue() > 0) { |
| assertEquals(numRecordsWritten * sizePerRecordWithOverhead, |
| outputBytesWithOverheadCounter.getValue()); |
| } else { |
| assertEquals(0, outputBytesWithOverheadCounter.getValue()); |
| } |
| long fileOutputBytes = fileOutputBytesCounter.getValue(); |
| if (numRecordsWritten > 0) { |
| assertTrue(fileOutputBytes > 0); |
| if (!shouldCompress) { |
| assertTrue("fileOutputBytes=" + fileOutputBytes + ", outputRecordBytes=" |
| +outputRecordBytesCounter.getValue(), |
| fileOutputBytes > outputRecordBytesCounter.getValue()); |
| } |
| } else { |
| assertEquals(0, fileOutputBytes); |
| } |
| // due to multiple threads, buffers could be merged in chunks in scheduleSpill. |
| assertTrue(recordsPerBuffer * numExpectedSpills >= spilledRecordsCounter.getValue()); |
| long additionalSpillBytesWritten = additionalSpillBytesWritternCounter.getValue(); |
| long additionalSpillBytesRead = additionalSpillBytesReadCounter.getValue(); |
| |
| //No additional spill bytes written when final merge is disabled. |
| assertEquals(additionalSpillBytesWritten, 0); |
| |
| //No additional spills when final merge is disabled. |
| assertTrue(additionalSpillBytesWritten == additionalSpillBytesRead); |
| |
| //No additional spills when final merge is disabled. |
| assertEquals(numAdditionalSpillsCounter.getValue(), 0); |
| |
| assertTrue(lastEvents.size() > 0); |
| //Get the last event |
| int index = lastEvents.size() - 1; |
| assertTrue(lastEvents.get(index) instanceof CompositeDataMovementEvent); |
| CompositeDataMovementEvent cdme = |
| (CompositeDataMovementEvent)lastEvents.get(index); |
| assertEquals(0, cdme.getSourceIndexStart()); |
| assertEquals(numOutputs, cdme.getCount()); |
| DataMovementEventPayloadProto eventProto = |
| DataMovementEventPayloadProto.parseFrom( |
| ByteString.copyFrom(cdme.getUserPayload())); |
| |
| verifyEmptyPartitions(eventProto, numRecordsWritten, numPartitions, skippedPartitions); |
| |
| if (outputRecordsCounter.getValue() > 0) { |
| //Ensure that this is the last event |
| assertTrue(eventProto.getLastEvent()); |
| } |
| |
| // Verify if all path components have spillIds when final merge is disabled |
| Pattern mergePathComponentPattern = Pattern.compile("(.*)(_\\d+)"); |
| for(Event event : lastEvents) { |
| if (!(event instanceof CompositeDataMovementEvent)) { |
| continue; |
| } |
| cdme = (CompositeDataMovementEvent)event; |
| eventProto = DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(cdme.getUserPayload())); |
| |
| assertEquals(false, eventProto.getPipelined()); |
| if (eventProto.hasPathComponent()) { |
| //for final merge disabled cases, it should have _spillId |
| Matcher matcher = mergePathComponentPattern.matcher(eventProto.getPathComponent()); |
| assertTrue("spill id should be present in path component " + eventProto.getPathComponent(), matcher.matches()); |
| assertEquals(2, matcher.groupCount()); |
| assertEquals(uniqueId, matcher.group(1)); |
| assertTrue("spill id should be present in path component", matcher.group(2) != null); |
| Path outputPath = new Path(outputContext.getWorkDirs()[0], |
| "output/" + eventProto.getPathComponent() + "/" + Constants.TEZ_RUNTIME_TASK_OUTPUT_FILENAME_STRING); |
| Path indexPath = outputPath.suffix(Constants.TEZ_RUNTIME_TASK_OUTPUT_INDEX_SUFFIX_STRING); |
| checkPermissions(outputPath, indexPath); |
| } else { |
| assertEquals(0, eventProto.getSpillId()); |
| if (outputRecordsCounter.getValue() > 0) { |
| assertEquals(true, eventProto.getLastEvent()); |
| } else { |
| byte[] emptyPartitions = TezCommonUtils.decompressByteStringToByteArray(eventProto |
| .getEmptyPartitions()); |
| BitSet emptyPartitionBits = TezUtilsInternal.fromByteArray(emptyPartitions); |
| assertEquals(numPartitions, emptyPartitionBits.cardinality()); |
| } |
| } |
| } |
| |
| |
| verify(outputContext, atLeast(1)).notifyProgress(); |
| |
| // Verify if all spill files are available. |
| TezTaskOutput taskOutput = new TezTaskOutputFiles(conf, uniqueId, dagId); |
| |
| if (numRecordsWritten > 0) { |
| int numSpills = kvWriter.numSpills.get(); |
| for (int i = 0; i < numSpills; i++) { |
| assertTrue(localFs.exists(taskOutput.getSpillFileForWrite(i, 10))); |
| assertTrue(localFs.exists(taskOutput.getSpillIndexFileForWrite(i, 10))); |
| } |
| } else { |
| return; |
| } |
| } |
| |
| private void baseTest(int numRecords, int numPartitions, Set<Integer> skippedPartitions, |
| boolean shouldCompress, int maxSingleBufferSizeBytes, int bufferMergePercent) |
| throws IOException, InterruptedException { |
| baseTest(numRecords, numPartitions, skippedPartitions, shouldCompress, |
| maxSingleBufferSizeBytes, bufferMergePercent, 2048, false); |
| } |
| |
| private void baseTest(int numRecords, int numPartitions, Set<Integer> skippedPartitions, |
| boolean shouldCompress, int maxSingleBufferSizeBytes, int bufferMergePercent, int |
| availableMemory, boolean dataViaEventEnabled) |
| throws IOException, InterruptedException { |
| PartitionerForTest partitioner = new PartitionerForTest(); |
| ApplicationId appId = ApplicationId.newInstance(10000000, 1); |
| TezCounters counters = new TezCounters(); |
| String uniqueId = UUID.randomUUID().toString(); |
| int dagId = 1; |
| String auxiliaryService = defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, |
| TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT); |
| OutputContext outputContext = createMockOutputContext(counters, appId, uniqueId, auxiliaryService); |
| |
| Configuration conf = createConfiguration(outputContext, IntWritable.class, LongWritable.class, |
| shouldCompress, maxSingleBufferSizeBytes); |
| conf.setInt( |
| TezRuntimeConfiguration.TEZ_RUNTIME_UNORDERED_PARTITIONED_KVWRITER_BUFFER_MERGE_PERCENT, |
| bufferMergePercent); |
| conf.setBoolean( |
| TezRuntimeConfiguration.TEZ_RUNTIME_TRANSFER_DATA_VIA_EVENTS_ENABLED, |
| dataViaEventEnabled); |
| |
| CompressionCodec codec = null; |
| if (shouldCompress) { |
| codec = new DefaultCodec(); |
| ((Configurable) codec).setConf(conf); |
| } |
| |
| int numOutputs = numPartitions; |
| int numRecordsWritten = 0; |
| |
| Map<Integer, Multimap<Integer, Long>> expectedValues = new HashMap<Integer, Multimap<Integer, Long>>(); |
| for (int i = 0; i < numOutputs; i++) { |
| expectedValues.put(i, LinkedListMultimap.<Integer, Long> create()); |
| } |
| |
| UnorderedPartitionedKVWriter kvWriter = new UnorderedPartitionedKVWriterForTest(outputContext, |
| conf, numOutputs, availableMemory); |
| |
| int sizePerBuffer = kvWriter.sizePerBuffer; |
| int sizePerRecord = 4 + 8; // IntW + LongW |
| int sizePerRecordWithOverhead = sizePerRecord + 12; // Record + META_OVERHEAD |
| |
| IntWritable intWritable = new IntWritable(); |
| LongWritable longWritable = new LongWritable(); |
| BitSet partitionsWithData = new BitSet(numPartitions); |
| for (int i = 0; i < numRecords; i++) { |
| intWritable.set(i); |
| longWritable.set(i); |
| int partition = partitioner.getPartition(intWritable, longWritable, numOutputs); |
| if (skippedPartitions != null && skippedPartitions.contains(partition)) { |
| continue; |
| } |
| partitionsWithData.set(partition); |
| expectedValues.get(partition).put(intWritable.get(), longWritable.get()); |
| kvWriter.write(intWritable, longWritable); |
| numRecordsWritten++; |
| } |
| List<Event> events = kvWriter.close(); |
| |
| if (numPartitions == 1) { |
| assertEquals(true, kvWriter.skipBuffers); |
| |
| // VM & DME events |
| assertEquals(2, events.size()); |
| Event event1 = events.get(1); |
| assertTrue(event1 instanceof CompositeDataMovementEvent); |
| CompositeDataMovementEvent dme = (CompositeDataMovementEvent) event1; |
| ByteBuffer bb = dme.getUserPayload(); |
| ShuffleUserPayloads.DataMovementEventPayloadProto shufflePayload = |
| ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(bb)); |
| assertEquals(kvWriter.outputRecordsCounter.getValue(), shufflePayload.getNumRecord()); |
| } |
| |
| int recordsPerBuffer = sizePerBuffer / sizePerRecordWithOverhead; |
| int numExpectedSpills = numRecordsWritten / recordsPerBuffer / kvWriter.spillLimit; |
| |
| verify(outputContext, never()).reportFailure(any(), any(), any()); |
| |
| assertNull(kvWriter.currentBuffer); |
| assertEquals(0, kvWriter.availableBuffers.size()); |
| |
| // Verify the counters |
| TezCounter outputRecordBytesCounter = counters.findCounter(TaskCounter.OUTPUT_BYTES); |
| TezCounter outputRecordsCounter = counters.findCounter(TaskCounter.OUTPUT_RECORDS); |
| TezCounter outputBytesWithOverheadCounter = counters |
| .findCounter(TaskCounter.OUTPUT_BYTES_WITH_OVERHEAD); |
| TezCounter fileOutputBytesCounter = counters.findCounter(TaskCounter.OUTPUT_BYTES_PHYSICAL); |
| TezCounter spilledRecordsCounter = counters.findCounter(TaskCounter.SPILLED_RECORDS); |
| TezCounter additionalSpillBytesWritternCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_WRITTEN); |
| TezCounter additionalSpillBytesReadCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILLS_BYTES_READ); |
| TezCounter numAdditionalSpillsCounter = counters |
| .findCounter(TaskCounter.ADDITIONAL_SPILL_COUNT); |
| assertEquals(numRecordsWritten * sizePerRecord, outputRecordBytesCounter.getValue()); |
| if (numPartitions > 1) { |
| assertEquals(numRecordsWritten * sizePerRecordWithOverhead, outputBytesWithOverheadCounter.getValue()); |
| } |
| assertEquals(numRecordsWritten, outputRecordsCounter.getValue()); |
| |
| long fileOutputBytes = fileOutputBytesCounter.getValue(); |
| if (numRecordsWritten > 0) { |
| assertTrue(fileOutputBytes > 0); |
| if ((!shouldCompress) && (!dataViaEventEnabled)) { |
| assertTrue(fileOutputBytes > outputRecordBytesCounter.getValue()); |
| } |
| } else { |
| assertEquals(0, fileOutputBytes); |
| } |
| if (!dataViaEventEnabled) { |
| assertEquals(recordsPerBuffer * numExpectedSpills, |
| spilledRecordsCounter.getValue()); |
| } |
| long additionalSpillBytesWritten = additionalSpillBytesWritternCounter.getValue(); |
| long additionalSpillBytesRead = additionalSpillBytesReadCounter.getValue(); |
| if (numExpectedSpills == 0) { |
| assertEquals(0, additionalSpillBytesWritten); |
| assertEquals(0, additionalSpillBytesRead); |
| } else { |
| assertTrue(additionalSpillBytesWritten > 0); |
| if (!dataViaEventEnabled) { |
| assertTrue(additionalSpillBytesRead > 0); |
| if (!shouldCompress) { |
| assertTrue(additionalSpillBytesWritten > |
| (recordsPerBuffer * numExpectedSpills * sizePerRecord)); |
| assertTrue(additionalSpillBytesRead > |
| (recordsPerBuffer * numExpectedSpills * sizePerRecord)); |
| } |
| } else { |
| if (kvWriter.writer.getCompressedLength() > |
| TezRuntimeConfiguration.TEZ_RUNTIME_TRANSFER_DATA_VIA_EVENTS_MAX_SIZE_DEFAULT) { |
| assertTrue(additionalSpillBytesWritten > 0); |
| } |
| } |
| } |
| if (!dataViaEventEnabled) { |
| assertEquals(additionalSpillBytesWritten, additionalSpillBytesRead); |
| } |
| // due to multiple threads, buffers could be merged in chunks in scheduleSpill. |
| assertTrue(numExpectedSpills >= numAdditionalSpillsCounter.getValue()); |
| |
| BitSet emptyPartitionBits = null; |
| // Verify the events returned |
| assertEquals(2, events.size()); |
| assertTrue(events.get(0) instanceof VertexManagerEvent); |
| VertexManagerEvent vme = (VertexManagerEvent) events.get(0); |
| verifyPartitionStats(vme, partitionsWithData); |
| assertTrue(events.get(1) instanceof CompositeDataMovementEvent); |
| CompositeDataMovementEvent cdme = (CompositeDataMovementEvent) events.get(1); |
| assertEquals(0, cdme.getSourceIndexStart()); |
| assertEquals(numOutputs, cdme.getCount()); |
| DataMovementEventPayloadProto eventProto = |
| DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( |
| cdme.getUserPayload())); |
| if (skippedPartitions == null && numRecordsWritten > 0) { |
| assertFalse(eventProto.hasEmptyPartitions()); |
| emptyPartitionBits = new BitSet(numPartitions); |
| } else { |
| assertTrue(eventProto.hasEmptyPartitions()); |
| byte[] emptyPartitions = TezCommonUtils.decompressByteStringToByteArray(eventProto |
| .getEmptyPartitions()); |
| emptyPartitionBits = TezUtilsInternal.fromByteArray(emptyPartitions); |
| if (numRecordsWritten == 0) { |
| assertEquals(numPartitions, emptyPartitionBits.cardinality()); |
| } else { |
| for (Integer e : skippedPartitions) { |
| assertTrue(emptyPartitionBits.get(e)); |
| } |
| assertEquals(skippedPartitions.size(), emptyPartitionBits.cardinality()); |
| } |
| } |
| if (emptyPartitionBits.cardinality() != numPartitions) { |
| assertEquals(HOST_STRING, eventProto.getHost()); |
| assertEquals(SHUFFLE_PORT, eventProto.getPort()); |
| assertEquals(uniqueId, eventProto.getPathComponent()); |
| } else { |
| assertFalse(eventProto.hasHost()); |
| assertFalse(eventProto.hasPort()); |
| assertFalse(eventProto.hasPathComponent()); |
| } |
| |
| // Verify the actual data |
| TezTaskOutput taskOutput = new TezTaskOutputFiles(conf, uniqueId, dagId); |
| Path outputFilePath = kvWriter.finalOutPath; |
| Path spillFilePath = kvWriter.finalIndexPath; |
| |
| if (numRecordsWritten <= 0) { |
| return; |
| } |
| |
| boolean isInMem= eventProto.getData().hasData(); |
| assertTrue(localFs.exists(outputFilePath)); |
| assertEquals("Incorrect output permissions (user)", FsAction.READ_WRITE, |
| localFs.getFileStatus(outputFilePath).getPermission().getUserAction()); |
| assertEquals("Incorrect output permissions (group)", FsAction.READ, |
| localFs.getFileStatus(outputFilePath).getPermission().getGroupAction()); |
| if( !isInMem ) { |
| assertTrue(localFs.exists(spillFilePath)); |
| assertEquals("Incorrect index permissions (user)", FsAction.READ_WRITE, |
| localFs.getFileStatus(spillFilePath).getPermission().getUserAction()); |
| assertEquals("Incorrect index permissions (group)", FsAction.READ, |
| localFs.getFileStatus(spillFilePath).getPermission().getGroupAction()); |
| |
| // verify no intermediate spill files have been left around |
| synchronized (kvWriter.spillInfoList) { |
| for (SpillInfo spill : kvWriter.spillInfoList) { |
| assertFalse("lingering intermediate spill file " + spill.outPath, |
| localFs.exists(spill.outPath)); |
| } |
| } |
| } |
| |
| // Special case for 0 records. |
| DataInputBuffer keyBuffer = new DataInputBuffer(); |
| DataInputBuffer valBuffer = new DataInputBuffer(); |
| IntWritable keyDeser = new IntWritable(); |
| LongWritable valDeser = new LongWritable(); |
| for (int i = 0; i < numOutputs; i++) { |
| IFile.Reader reader = null; |
| InputStream inStream; |
| if (isInMem) { |
| // Read from in memory payload |
| int dataLoadSize = eventProto.getData().getData().size(); |
| inStream = new ByteArrayInputStream(eventProto.getData().getData().toByteArray()); |
| reader = new IFile.Reader(inStream, dataLoadSize, codec, null, |
| null, false, 0, -1); |
| } else { |
| TezSpillRecord spillRecord = new TezSpillRecord(spillFilePath, conf); |
| TezIndexRecord indexRecord = spillRecord.getIndex(i); |
| if (skippedPartitions != null && skippedPartitions.contains(i)) { |
| assertFalse("The Index Record for partition " + i + " should not have any data", indexRecord.hasData()); |
| continue; |
| } |
| |
| FSDataInputStream tmpStream = FileSystem.getLocal(conf).open(outputFilePath); |
| tmpStream.seek(indexRecord.getStartOffset()); |
| inStream = tmpStream; |
| reader = new IFile.Reader(tmpStream, indexRecord.getPartLength(), codec, null, |
| null, false, 0, -1); |
| } |
| |
| while (reader.nextRawKey(keyBuffer)) { |
| reader.nextRawValue(valBuffer); |
| keyDeser.readFields(keyBuffer); |
| valDeser.readFields(valBuffer); |
| int partition = partitioner.getPartition(keyDeser, valDeser, numOutputs); |
| assertTrue(expectedValues.get(partition).remove(keyDeser.get(), valDeser.get())); |
| } |
| inStream.close(); |
| } |
| for (int i = 0; i < numOutputs; i++) { |
| assertEquals(0, expectedValues.get(i).size()); |
| expectedValues.remove(i); |
| } |
| assertEquals(0, expectedValues.size()); |
| verify(outputContext, atLeast(1)).notifyProgress(); |
| } |
| |
| private static String createRandomString(int size) { |
| StringBuilder sb = new StringBuilder(size); |
| Random random = new Random(); |
| for (int i = 0; i < size; i++) { |
| int r = Math.abs(random.nextInt()) % 26; |
| sb.append((char) (65 + r)); |
| } |
| return sb.toString(); |
| } |
| |
| private OutputContext createMockOutputContext(TezCounters counters, ApplicationId appId, |
| String uniqueId, String auxiliaryService) { |
| OutputContext outputContext = mock(OutputContext.class); |
| doReturn(counters).when(outputContext).getCounters(); |
| doReturn(appId).when(outputContext).getApplicationId(); |
| doReturn(1).when(outputContext).getDAGAttemptNumber(); |
| doReturn("dagName").when(outputContext).getDAGName(); |
| doReturn("destinationVertexName").when(outputContext).getDestinationVertexName(); |
| doReturn(1).when(outputContext).getOutputIndex(); |
| doReturn(1).when(outputContext).getTaskAttemptNumber(); |
| doReturn(1).when(outputContext).getTaskIndex(); |
| doReturn(1).when(outputContext).getTaskVertexIndex(); |
| doReturn("vertexName").when(outputContext).getTaskVertexName(); |
| doReturn(uniqueId).when(outputContext).getUniqueIdentifier(); |
| |
| doAnswer(new Answer<ByteBuffer>() { |
| @Override |
| public ByteBuffer answer(InvocationOnMock invocation) throws Throwable { |
| ByteBuffer portBuffer = ByteBuffer.allocate(4); |
| portBuffer.mark(); |
| portBuffer.putInt(SHUFFLE_PORT); |
| portBuffer.reset(); |
| return portBuffer; |
| } |
| }).when(outputContext).getServiceProviderMetaData(eq(auxiliaryService)); |
| |
| Path outDirBase = new Path(TEST_ROOT_DIR, "outDir_" + uniqueId); |
| String[] outDirs = new String[] { outDirBase.toString() }; |
| doReturn(outDirs).when(outputContext).getWorkDirs(); |
| return outputContext; |
| } |
| |
| private Configuration createConfiguration(OutputContext outputContext, |
| Class<? extends Writable> keyClass, Class<? extends Writable> valClass, |
| boolean shouldCompress, int maxSingleBufferSizeBytes) { |
| return createConfiguration(outputContext, keyClass, valClass, shouldCompress, |
| maxSingleBufferSizeBytes, PartitionerForTest.class); |
| } |
| |
| private Configuration createConfiguration(OutputContext outputContext, |
| Class<? extends Writable> keyClass, Class<? extends Writable> valClass, |
| boolean shouldCompress, int maxSingleBufferSizeBytes, |
| Class<? extends Partitioner> partitionerClass) { |
| Configuration conf = new Configuration(false); |
| conf.set(CommonConfigurationKeys.FS_PERMISSIONS_UMASK_KEY, "077"); |
| conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, outputContext.getWorkDirs()); |
| conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS, keyClass.getName()); |
| conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS, valClass.getName()); |
| conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS, partitionerClass.getName()); |
| if (maxSingleBufferSizeBytes >= 0) { |
| conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_UNORDERED_OUTPUT_MAX_PER_BUFFER_SIZE_BYTES, |
| maxSingleBufferSizeBytes); |
| } |
| conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_COMPRESS, shouldCompress); |
| if (shouldCompress) { |
| conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_COMPRESS_CODEC, |
| DefaultCodec.class.getName()); |
| } |
| conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS, |
| reportPartitionStats.getType()); |
| return conf; |
| } |
| |
| public static class PartitionerForTest implements Partitioner { |
| @Override |
| public int getPartition(Object key, Object value, int numPartitions) { |
| if (key instanceof IntWritable) { |
| return ((IntWritable) key).get() % numPartitions; |
| } else { |
| throw new UnsupportedOperationException( |
| "Test partitioner expected to be called with IntWritable only"); |
| } |
| } |
| } |
| |
| private static class UnorderedPartitionedKVWriterForTest extends UnorderedPartitionedKVWriter { |
| |
| public UnorderedPartitionedKVWriterForTest(OutputContext outputContext, Configuration conf, |
| int numOutputs, long availableMemoryBytes) throws IOException { |
| super(outputContext, conf, numOutputs, availableMemoryBytes); |
| } |
| |
| @Override |
| String getHost() { |
| return HOST_STRING; |
| } |
| } |
| } |
