diff --git a/.codespellrc b/.codespellrc index fee7e92ef2..6c3fdddf2d 100644 --- a/.codespellrc +++ b/.codespellrc @@ -18,4 +18,4 @@ # [codespell] -ignore-words-list = thirdparty,afterall,AfterAll +ignore-words-list = thirdparty,afterall,AfterAll,atleast,atLeast diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/ShuffleUtils.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/ShuffleUtils.java index b2135b0566..89468f4369 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/ShuffleUtils.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/ShuffleUtils.java @@ -337,6 +337,18 @@ static ByteBuffer generateDMEPayload(boolean sendEmptyPartitionDetails, return payload; } + /** + * Fills in num_record, which describes a single input, so only a one-partition payload can + * carry it: with more partitions the one payload covers several inputs and the count could not + * be attributed to any of them. + */ + public static void setNumRecord(DataMovementEventPayloadProto.Builder payloadBuilder, + int numPartitions, long numRecords) { + if (numPartitions == 1) { + payloadBuilder.setNumRecord(numRecords); + } + } + /** * Generate events for outputs which have not been started. * @param eventList @@ -356,6 +368,8 @@ public static void generateEventsForNonStartedOutput(List eventList, DataMovementEventPayloadProto.Builder payloadBuilder = DataMovementEventPayloadProto .newBuilder(); + // This output produced nothing; 0 says so, where an unset field reads as unknown. + setNumRecord(payloadBuilder, numPhysicalOutputs, 0); // Construct the VertexManager event if required. if (generateVmEvent) { diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleInputEventHandlerImpl.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleInputEventHandlerImpl.java index b832e97720..15f5f7c0e5 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleInputEventHandlerImpl.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleInputEventHandlerImpl.java @@ -176,20 +176,16 @@ private void processDataMovementEvent(DataMovementEvent dme, DataMovementEventPa .stringify(shufflePayload)); } - if (shufflePayload.hasEmptyPartitions()) { - if (emptyPartitionsBitSet.get(srcIndex)) { - CompositeInputAttemptIdentifier srcAttemptIdentifier = - constructInputAttemptIdentifier(dme.getTargetIndex(), 1, dme.getVersion(), shufflePayload, false); - LOG.debug("Source partition: {} did not generate any data. SrcAttempt: [{}]. Not fetching.", + updateApproximateInputRecords(dme.getTargetIndex(), shufflePayload); + + if (shufflePayload.hasEmptyPartitions() && emptyPartitionsBitSet.get(srcIndex)) { + CompositeInputAttemptIdentifier srcAttemptIdentifier = + constructInputAttemptIdentifier(dme.getTargetIndex(), 1, dme.getVersion(), shufflePayload, false); + LOG.debug("Source partition: {} did not generate any data. SrcAttempt: [{}]. Not fetching.", srcIndex, srcAttemptIdentifier); - numDmeEventsNoData.getAndIncrement(); - shuffleManager.addCompletedInputWithNoData(srcAttemptIdentifier.expand(0)); - return; - } else { - shuffleManager.updateApproximateInputRecords(shufflePayload.getNumRecord()); - } - } else { - shuffleManager.updateApproximateInputRecords(shufflePayload.getNumRecord()); + numDmeEventsNoData.getAndIncrement(); + shuffleManager.addCompletedInputWithNoData(srcAttemptIdentifier.expand(0)); + return; } CompositeInputAttemptIdentifier srcAttemptIdentifier = constructInputAttemptIdentifier(dme.getTargetIndex(), 1, dme.getVersion(), @@ -225,8 +221,23 @@ private void moveDataToFetchedInput(DataProto dataProto, } } + /** + * Tells the ShuffleManager how many rows one input has reported. Only a single-partition writer + * sets numRecord, so the event covers exactly one input -- a composite event's count is 1 -- and + * its target index is that input's. It is set even when the input wrote nothing, which is what + * lets an empty input count towards the extrapolation instead of being scaled over. + */ + private void updateApproximateInputRecords(int targetIndex, + DataMovementEventPayloadProto shufflePayload) { + if (shufflePayload.hasNumRecord()) { + shuffleManager.updateApproximateInputRecords(targetIndex, shufflePayload.getNumRecord()); + } + } + private void processCompositeRoutedDataMovementEvent(CompositeRoutedDataMovementEvent crdme, DataMovementEventPayloadProto shufflePayload, BitSet emptyPartitionsBitSet) throws IOException { int partitionId = crdme.getSourceIndex(); + updateApproximateInputRecords(crdme.getTargetIndex(), shufflePayload); + if (LOG.isDebugEnabled()) { LOG.debug("DME srcIdx: " + partitionId + ", targetIndex: " + crdme.getTargetIndex() + ", count:" + crdme.getCount() + ", attemptNum: " + crdme.getVersion() + ", payload: " + ShuffleUtils diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleManager.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleManager.java index 666c7b9ebb..0c0aed2215 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleManager.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleManager.java @@ -192,8 +192,10 @@ public class ShuffleManager implements FetcherCallback { private final AtomicBoolean isShutdown = new AtomicBoolean(false); - private long inputRecordsFromEvents; - private long eventsReceived; + private final long[] numRecordsPerInput; + private final BitSet reportedInputSet; + private long inputRecordsReported; + private final TezCounter approximateInputRecords; private final TezCounter shuffledInputsCounter; private final TezCounter failedShufflesCounter; @@ -230,6 +232,8 @@ public ShuffleManager(InputContext inputContext, Configuration conf, int numInpu this.numInputs = numInputs; this.approximateInputRecords = inputContext.getCounters().findCounter(TaskCounter.APPROXIMATE_INPUT_RECORDS); + this.numRecordsPerInput = new long[numInputs]; + this.reportedInputSet = new BitSet(numInputs); this.shuffledInputsCounter = inputContext.getCounters().findCounter(TaskCounter.NUM_SHUFFLED_INPUTS); this.failedShufflesCounter = inputContext.getCounters().findCounter(TaskCounter.NUM_FAILED_SHUFFLE_INPUTS); this.bytesShuffledCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_BYTES); @@ -348,13 +352,40 @@ public ShuffleManager(InputContext inputContext, Configuration conf, int numInpu + ", asyncHttp=" + asyncHttp); } - public void updateApproximateInputRecords(int delta) { - if (delta <= 0) { + /** + * Records how many rows one input reports and republishes the estimate over every input. + * A report is the source attempt's output record count so far, not a per-event delta, so a + * pipelined input reports a growing total as it spills and only its largest report counts. An + * input that wrote no rows reports zero and still joins the denominator, which numInputs + * already counts in the multiplier. + *

+ * The reports carry no attempt number, so a retried attempt that produces fewer rows than the + * one it replaces leaves the larger total in place. Monotonic per input is the only rule that + * composes across a pipelined input's spills without also subtracting on a reordered report, + * and the result is an estimate over the inputs that have reported -- never a bound on what + * the consumer will read. + */ + void updateApproximateInputRecords(int inputIndex, long numRecords) { + long lastReported = numRecordsPerInput[inputIndex]; + if (!reportedInputSet.get(inputIndex)) { + reportedInputSet.set(inputIndex); + } else if (numRecords <= lastReported) { return; } - inputRecordsFromEvents += delta; - eventsReceived++; - approximateInputRecords.setValue((inputRecordsFromEvents / eventsReceived) * numInputs); + inputRecordsReported += numRecords - lastReported; + numRecordsPerInput[inputIndex] = numRecords; + approximateInputRecords.setValue( + extrapolateTotal(inputRecordsReported, reportedInputSet.cardinality(), numInputs)); + } + + /** + * The mean over the inputs that reported, applied to all of them. The remainder is scaled as + * well, so the truncation is not multiplied by every input, and recordsReported * allInputs is + * never formed, so only the result has to fit in a long. + */ + private static long extrapolateTotal(long recordsReported, int reportedInputs, int allInputs) { + return (recordsReported / reportedInputs) * allInputs + + ((recordsReported % reportedInputs) * allInputs) / reportedInputs; } public void run() throws IOException { diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java index e644d8c12e..7b5e8f922b 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java @@ -861,13 +861,16 @@ public List close() throws IOException, InterruptedException { } //For pipelined case, send out an event in case finalspill generated a spill file. - if (finalSpill() != null) { + SpillResult result = finalSpill(); + // The events below read the task counters, and this buffer's rows reach them only here. + // Not before finalSpill: its empty-buffer branch reads localOutputRecordsCounter. + updateTezCountersAndNotify(); + if (result != null) { // VertexManagerEvent is only sent at the end and thus sizePerPartition is used // for the sum of all spills. mayBeSendEventsForSpill(currentBuffer.recordsPerPartition, sizePerPartition, numSpills.get() - 1, true); } - updateTezCountersAndNotify(); cleanupCurrentBuffer(); return events; } @@ -905,9 +908,10 @@ private Event generateDMEvent(boolean addSpillDetails, int spillId, outputContext.notifyProgress(); DataMovementEventPayloadProto.Builder payloadBuilder = DataMovementEventPayloadProto .newBuilder(); - if (numPartitions == 1) { - payloadBuilder.setNumRecord((int) outputRecordsCounter.getValue()); - } + // writeLargeRecord bypasses outputRecordsCounter, so the count has to add them back. This is + // the sum ShuffleUtils.generateVMEvent reports. + ShuffleUtils.setNumRecord(payloadBuilder, numPartitions, + outputRecordsCounter.getValue() + outputLargeRecordsCounter.getValue()); String host = getHost(); if (emptyPartitions.cardinality() != 0) { diff --git a/tez-runtime-library/src/main/proto/ShufflePayloads.proto b/tez-runtime-library/src/main/proto/ShufflePayloads.proto index 5cbd18a9b0..8e6c6c0492 100644 --- a/tez-runtime-library/src/main/proto/ShufflePayloads.proto +++ b/tez-runtime-library/src/main/proto/ShufflePayloads.proto @@ -30,7 +30,7 @@ message DataMovementEventPayloadProto { optional bool pipelined = 7; // Related to pipelined shuffle optional bool last_event = 8; // Related to pipelined shuffle optional int32 spill_id = 9; // Related to pipelined shuffle. - optional int32 num_record = 10; + optional int64 num_record = 10; } message DataProto { diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestShuffleUtils.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestShuffleUtils.java index 1891cf4609..54fe59e120 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestShuffleUtils.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestShuffleUtils.java @@ -286,6 +286,36 @@ public void testGenerateOnSpillEvent_With_All_EmptyPartitions() throws Exception "emptyPartitionBitSet cardinality (expecting 10) = " + emptyPartitionsBitSet.cardinality()); } + /** + * An output that never started delivers no rows, and a single-partition one fills num_record + * like any other, so it joins the consumer's denominator instead of being scaled over. + */ + @Test + public void testNonStartedSinglePartitionOutputReportsZeroRecords() throws Exception { + List events = Lists.newLinkedList(); + ShuffleUtils.generateEventsForNonStartedOutput(events, 1, outputContext, false, true, + TezCommonUtils.newBestCompressionDeflater()); + + ShuffleUserPayloads.DataMovementEventPayloadProto proto = + ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( + ((CompositeDataMovementEvent) events.get(0)).getUserPayload())); + assertTrue(proto.hasNumRecord()); + assertEquals(0, proto.getNumRecord()); + } + + /** num_record describes a single input, so a multi-partition payload must not carry one. */ + @Test + public void testNonStartedMultiPartitionOutputReportsNoRecordCount() throws Exception { + List events = Lists.newLinkedList(); + ShuffleUtils.generateEventsForNonStartedOutput(events, 10, outputContext, false, true, + TezCommonUtils.newBestCompressionDeflater()); + + ShuffleUserPayloads.DataMovementEventPayloadProto proto = + ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( + ((CompositeDataMovementEvent) events.get(0)).getUserPayload())); + assertFalse(proto.hasNumRecord()); + } + @Test public void testInternalErrorTranslation() throws Exception { String codecErrorMsg = "codec failure"; diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleInputEventHandlerImpl.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleInputEventHandlerImpl.java index 0eed8c65a3..c49b7f74d5 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleInputEventHandlerImpl.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleInputEventHandlerImpl.java @@ -408,6 +408,63 @@ public void testDataMovementEventsWithShuffleData() throws IOException { } } + /** + * An input whose partition is empty reports zero rows and must still reach the ShuffleManager: + * counting it only in numInputs and not in the denominator is what inflated the counter. + */ + @Test + @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS) + public void testAnEmptyPartitionStillReportsItsRecordCount() throws IOException { + ShuffleManager shuffleManager = mock(ShuffleManager.class); + ShuffleInputEventHandlerImpl handler = newHandler(shuffleManager, false); + + handler.handleEvents(Collections.singletonList( + createRecordCountEvent(1, 0, createEmptyPartitionByteString(0), false))); + + verify(shuffleManager).updateApproximateInputRecords(eq(1), eq(0L)); + verify(shuffleManager).addCompletedInputWithNoData(any()); + } + + /** The composite path carries the same payload and has to report it too. */ + @Test + @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS) + public void testACompositeEventReportsItsRecordCount() throws IOException { + ShuffleManager shuffleManager = mock(ShuffleManager.class); + ShuffleInputEventHandlerImpl handler = newHandler(shuffleManager, true); + + handler.handleEvents(Collections.singletonList(createRecordCountEvent(1, 4000, null, true))); + + verify(shuffleManager).updateApproximateInputRecords(eq(1), eq(4000L)); + } + + /** + * A source task can emit more rows than an int32 holds. num_record is an int64 so the count + * survives the wire; passing it as an int would not even compile against that field. + */ + @Test + @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS) + public void testACountAboveIntMaxSurvivesTheWire() throws IOException { + ShuffleManager shuffleManager = mock(ShuffleManager.class); + ShuffleInputEventHandlerImpl handler = newHandler(shuffleManager, false); + + handler.handleEvents(Collections.singletonList( + createRecordCountEvent(1, 3_000_000_000L, null, false))); + + verify(shuffleManager).updateApproximateInputRecords(eq(1), eq(3_000_000_000L)); + } + + /** A payload with no record count -- any multi-partition writer -- reports nothing. */ + @Test + @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS) + public void testAPayloadWithoutARecordCountReportsNothing() throws IOException { + ShuffleManager shuffleManager = mock(ShuffleManager.class); + ShuffleInputEventHandlerImpl handler = newHandler(shuffleManager, false); + + handler.handleEvents(Collections.singletonList(createDataMovementEvent(0, 1, null))); + + verify(shuffleManager, times(0)).updateApproximateInputRecords(anyInt(), anyLong()); + } + private Event createDataMovementEvent(boolean addSpillDetails, int srcIdx, int targetIdx, int spillId, boolean isLastSpill, BitSet emptyPartitions, int numPartitions, int attemptNum) throws IOException { @@ -438,7 +495,7 @@ private Event createDataMovementEvent(boolean addSpillDetails, int srcIdx, int t return DataMovementEvent.create(srcIdx, targetIdx, attemptNum, payload); } - private Event createDataMovementEvent(int srcIndex, int targetIndex, + private DataMovementEventPayloadProto.Builder createPayloadBuilder( ByteString emptyPartitionByteString) { DataMovementEventPayloadProto.Builder builder = DataMovementEventPayloadProto.newBuilder(); builder.setHost(HOST); @@ -447,9 +504,29 @@ private Event createDataMovementEvent(int srcIndex, int targetIndex, if (emptyPartitionByteString != null) { builder.setEmptyPartitions(emptyPartitionByteString); } - Event dme = DataMovementEvent - .create(srcIndex, targetIndex, 0, builder.build().toByteString().asReadOnlyByteBuffer()); - return dme; + return builder; + } + + private Event createDataMovementEvent(int srcIndex, int targetIndex, + ByteString emptyPartitionByteString) { + return DataMovementEvent.create(srcIndex, targetIndex, 0, + createPayloadBuilder(emptyPartitionByteString).build().toByteString() + .asReadOnlyByteBuffer()); + } + + private Event createRecordCountEvent(int targetIndex, long numRecord, + ByteString emptyPartitionByteString, boolean composite) { + ByteBuffer payload = createPayloadBuilder(emptyPartitionByteString).setNumRecord(numRecord) + .build().toByteString().asReadOnlyByteBuffer(); + return composite + ? CompositeRoutedDataMovementEvent.create(0, targetIndex, 1, 0, payload) + : DataMovementEvent.create(0, targetIndex, 0, payload); + } + + private ShuffleInputEventHandlerImpl newHandler(ShuffleManager shuffleManager, + boolean compositeFetch) { + return new ShuffleInputEventHandlerImpl(mock(InputContext.class), shuffleManager, + mock(FetchedInputAllocator.class), null, false, 0, compositeFetch); } private ByteString createEmptyPartitionByteString(int... emptyPartitions) throws IOException { diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleManager.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleManager.java index 90d8a551d1..20d4587cfc 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleManager.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/impl/TestShuffleManager.java @@ -49,6 +49,7 @@ import org.apache.tez.common.TezExecutors; import org.apache.tez.common.TezRuntimeFrameworkConfigs; import org.apache.tez.common.TezSharedExecutor; +import org.apache.tez.common.counters.TaskCounter; import org.apache.tez.common.counters.TezCounters; import org.apache.tez.common.security.JobTokenIdentifier; import org.apache.tez.common.security.JobTokenSecretManager; @@ -160,6 +161,94 @@ public void testMultiplePartitions() throws Exception { shuffleManager.getNumOfCompletedInputs()); } + /** + * An input that wrote no rows reports zero. Counting it only in numInputs and not in the + * denominator spreads the mean of the inputs that had data over the ones that did not. + */ + @Test + public void testInputsThatWroteNothingJoinTheDenominator() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 10); + shuffleManager.updateApproximateInputRecords(0, 1000); + shuffleManager.updateApproximateInputRecords(1, 1000); + for (int input = 2; input < 10; input++) { + shuffleManager.updateApproximateInputRecords(input, 0); + } + assertEquals(2000L, approximateInputRecords(inputContext)); + } + + /** + * numRecord is the source task's running OUTPUT_RECORDS, so a pipelined input reports a new + * total per spill. Summing the reports counts the earlier spills again. + */ + @Test + public void testOnlyTheLatestReportPerInputIsCounted() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 2); + shuffleManager.updateApproximateInputRecords(0, 100); + shuffleManager.updateApproximateInputRecords(0, 200); + shuffleManager.updateApproximateInputRecords(0, 300); + shuffleManager.updateApproximateInputRecords(1, 300); + assertEquals(600L, approximateInputRecords(inputContext)); + } + + /** Read before every input has reported, the value is the mean so far across all of them. */ + @Test + public void testExtrapolatesOverTheInputsThatHaveNotReported() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 8); + shuffleManager.updateApproximateInputRecords(0, 500); + assertEquals(4000L, approximateInputRecords(inputContext)); + shuffleManager.updateApproximateInputRecords(1, 300); + assertEquals(3200L, approximateInputRecords(inputContext)); + } + + /** + * The remainder of the mean is scaled too. Scaling only the quotient multiplies its truncation + * by every input: 301 rows over 3 reporting inputs of 10 reads 1000, not 1003. + */ + @Test + public void testTheMeanIsNotTruncatedBeforeItIsScaled() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 10); + shuffleManager.updateApproximateInputRecords(0, 100); + shuffleManager.updateApproximateInputRecords(1, 100); + shuffleManager.updateApproximateInputRecords(2, 101); + assertEquals(1003L, approximateInputRecords(inputContext)); + } + + /** + * The mean is scaled without ever forming recordsReported * numInputs, so only the answer has + * to fit in a long: 2e18 rows over 2 of 8 inputs is 8e18, while multiplying first would be + * 1.6e19 and wrap negative. + */ + @Test + public void testOnlyTheResultHasToFitInALong() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 8); + shuffleManager.updateApproximateInputRecords(0, 1_000_000_000_000_000_000L); + shuffleManager.updateApproximateInputRecords(1, 1_000_000_000_000_000_000L); + assertEquals(8_000_000_000_000_000_000L, approximateInputRecords(inputContext)); + } + + /** + * Spill callbacks read the record counter and send their event without a lock between, so a + * smaller total can arrive after a larger one. Subtracting it would drop the counter by the + * difference times the extrapolation factor. + */ + @Test + public void testAReportThatLostTheRaceIsIgnored() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 2); + shuffleManager.updateApproximateInputRecords(0, 1000); + shuffleManager.updateApproximateInputRecords(0, 500); + assertEquals(2000L, approximateInputRecords(inputContext)); + } + + private long approximateInputRecords(InputContext inputContext) { + return inputContext.getCounters().findCounter(TaskCounter.APPROXIMATE_INPUT_RECORDS).getValue(); + } + private InputContext createInputContext() throws IOException { DataOutputBuffer port_dob = new DataOutputBuffer(); port_dob.writeInt(PORT); diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java index 82fef1a886..29991818c1 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java @@ -41,6 +41,7 @@ import java.io.IOException; import java.io.InputStream; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.BitSet; import java.util.HashMap; import java.util.Iterator; @@ -432,6 +433,72 @@ public void testTextMixedRecordsWithoutFinalMerge(boolean shouldCompress, Report textTest(100, 10, 2048, 10, 10, 10, false, false); } + /** + * One partition with pipelined shuffle is the only shape where writeLargeRecord runs while + * num_record is set, so it is the only one that can show the DME under-counting large records. + */ + @ParameterizedTest(name = "test[{0}, {1}]") + @MethodSource("data") + @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS) + public void testLargeRecords_SinglePartition(boolean shouldCompress, + ReportPartitionStats reportPartitionStats) throws IOException, InterruptedException { + setupInit(shouldCompress, reportPartitionStats); + textTest(0, 1, 2048, 0, 0, 5, true, false); + } + + /** + * An output small enough to fit one buffer spills only at close, so its records reach the task + * counters in the same call that builds the final event -- which is where a count read too + * early reports zero. + */ + @ParameterizedTest(name = "test[{0}, {1}]") + @MethodSource("data") + @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS) + public void testFinalEventRecordCount_SinglePartition(boolean shouldCompress, + ReportPartitionStats reportPartitionStats) throws IOException, InterruptedException { + setupInit(shouldCompress, reportPartitionStats); + textTest(5, 1, 2048, 0, 0, 0, true, false); + } + + /** + * A single-partition writer is the only one that fills num_record, and large records never + * reach outputRecordsCounter -- writeLargeRecord bypasses it -- so both the DME and the + * VertexManager event have to carry the sum of the two counters. Under pipelined shuffle the + * events are pushed through sendEvents rather than returned by close(), so both sources are + * scanned and the last of each kind is the one that has to be right. + */ + private void assertLastEventsCountEveryRecord(UnorderedPartitionedKVWriter kvWriter, + OutputContext outputContext, List closeEvents) throws IOException { + List allEvents = new ArrayList<>(closeEvents); + @SuppressWarnings("unchecked") + ArgumentCaptor> sent = ArgumentCaptor.forClass(List.class); + verify(outputContext, atLeast(0)).sendEvents(sent.capture()); + sent.getAllValues().forEach(allEvents::addAll); + + long expected = kvWriter.outputRecordsCounter.getValue() + + kvWriter.outputLargeRecordsCounter.getValue(); + long lastDmeCount = -1; + long lastVmCount = -1; + for (Event event : allEvents) { + if (event instanceof CompositeDataMovementEvent) { + ShuffleUserPayloads.DataMovementEventPayloadProto dme = + ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( + ((CompositeDataMovementEvent) event).getUserPayload())); + if (dme.hasNumRecord()) { + lastDmeCount = dme.getNumRecord(); + } + } else if (event instanceof VertexManagerEvent) { + ShuffleUserPayloads.VertexManagerEventPayloadProto vme = + ShuffleUserPayloads.VertexManagerEventPayloadProto.parseFrom(ByteString.copyFrom( + ((VertexManagerEvent) event).getUserPayload())); + lastVmCount = vme.getNumRecord(); + } + } + assertTrue(lastDmeCount >= 0, "expected a DME carrying num_record"); + assertEquals(expected, lastDmeCount, "DME"); + assertEquals(expected, lastVmCount, "VertexManagerEvent"); + } + public void textTest(int numRegularRecords, int numPartitions, long availableMemory, int numLargeKeys, int numLargevalues, int numLargeKvPairs, boolean pipeliningEnabled, boolean isFinalMergeEnabled) throws IOException, @@ -536,6 +603,10 @@ public void textTest(int numRegularRecords, int numPartitions, long availableMem List events = kvWriter.close(); verify(outputContext, never()).reportFailure(any(), any(), any()); + if (numPartitions == 1) { + assertLastEventsCountEveryRecord(kvWriter, outputContext, events); + } + if (!pipeliningEnabled) { VertexManagerEvent vmEvent = null; for (Event event : events) { @@ -1321,7 +1392,12 @@ private void baseTest(int numRecords, int numPartitions, Set skippedPar ByteBuffer bb = dme.getUserPayload(); ShuffleUserPayloads.DataMovementEventPayloadProto shufflePayload = ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(bb)); - assertEquals(kvWriter.outputRecordsCounter.getValue(), shufflePayload.getNumRecord()); + // Large records bypass outputRecordsCounter, so the DME carries the same sum the + // VertexManager event does. They are zero on this path (skipBuffers), so this pins the + // common case only -- writeLargeRecord needs numPartitions == 1 with pipelined shuffle. + assertEquals(kvWriter.outputRecordsCounter.getValue() + + kvWriter.outputLargeRecordsCounter.getValue(), + shufflePayload.getNumRecord()); } int recordsPerBuffer = sizePerBuffer / sizePerRecordWithOverhead;