From 9d9a96ce77fdff607a85bf1cb0b123c2868f873a Mon Sep 17 00:00:00 2001 From: Denys Kuzmenko Date: Sat, 19 Sep 2026 13:44:34 +0300 Subject: [PATCH] TEZ-4757: APPROXIMATE_INPUT_RECORDS over-counts the rows an input will deliver - A report carries the source task's output record count so far, not a per-event delta; keep the latest value per input instead of summing events, so a pipelined input's earlier spills are not counted again - Ignore a report smaller than that input's last one: spill callbacks read the counter and send their event with no lock between, so a stale total can arrive last and would otherwise subtract from the sum - Count an input that wrote no rows in the denominator; leaving it out while numInputs keeps it in the multiplier spread the mean of the inputs that had data over the ones that did not - An output that never started reports 0 rather than nothing, so it joins the denominator instead of being scaled over - Flush the record counters before the final pipelined events read them. An output small enough for one buffer spills only at close, so its DME and its VertexManager event both carried 0 - Report from the composite event path too: which path runs depends only on tez.am.shuffle.auxiliary-service.id, so on the Tez shuffle handler the counter was never updated at all - Scale the mean without forming sum * numInputs, so only the true answer has to fit in a long - num_record is int64, matching the VertexManagerEventPayloadProto field it mirrors. It was int32 while both are filled from long counters, so a task emitting more than 2^31 rows wrapped negative; widening the field alone would only move the loss to the consumer, so the per-input array and the update method carry a long as well. int32 -> int64 is wire compatible in both directions on the same field number - The count includes OUTPUT_LARGE_RECORDS, which bypass OUTPUT_RECORDS; this is what ShuffleUtils.generateVMEvent already sends. Without it a single-partition writer reports 0 for a broadcast of large rows - updateApproximateInputRecords is package-private: it is shuffle internals with one caller in the same package - Measured before the fix: 10 inputs, 2 reporting 1000 rows each, read 10,000 instead of 2,000; 4 pipelined events totalling 600 read 450; a 5-row single-partition pipelined output reported 0 --- .codespellrc | 2 +- .../library/common/shuffle/ShuffleUtils.java | 14 +++ .../impl/ShuffleInputEventHandlerImpl.java | 37 +++++--- .../common/shuffle/impl/ShuffleManager.java | 45 ++++++++-- .../writers/UnorderedPartitionedKVWriter.java | 14 +-- .../src/main/proto/ShufflePayloads.proto | 2 +- .../common/shuffle/TestShuffleUtils.java | 30 +++++++ .../TestShuffleInputEventHandlerImpl.java | 85 +++++++++++++++++- .../shuffle/impl/TestShuffleManager.java | 89 +++++++++++++++++++ .../TestUnorderedPartitionedKVWriter.java | 78 +++++++++++++++- 10 files changed, 364 insertions(+), 32 deletions(-) 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;