From 24d7eebe84633016c7341234ef8b9be2bebc1637 Mon Sep 17 00:00:00 2001 From: Denys Kuzmenko Date: Sat, 19 Sep 2026 13:44:34 +0300 Subject: [PATCH 1/2] 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; From e22cc8a38eb31ed9d198d6f7827e1b6b5a149ea9 Mon Sep 17 00:00:00 2001 From: Denys Kuzmenko Date: Wed, 30 Sep 2026 15:57:54 +0300 Subject: [PATCH 2/2] TEZ-4757: review round - A negative report is rejected: it neither enters the sum nor marks the input as reported - Writer tests assert against the records the test wrote, with messages naming the event and counter that disagree; the no-data single-partition pipelined case is pinned in the harness that already runs it - Underscore-free test names, javadocs state what is pinned, helper below the tests, imported proto type used unqualified --- .../common/shuffle/impl/ShuffleManager.java | 17 +++++----- .../common/shuffle/TestShuffleUtils.java | 11 +++---- .../TestShuffleInputEventHandlerImpl.java | 5 ++- .../shuffle/impl/TestShuffleManager.java | 17 ++++++++-- .../TestUnorderedPartitionedKVWriter.java | 31 ++++++++++++------- 5 files changed, 53 insertions(+), 28 deletions(-) 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 0c0aed2215..48897628cf 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 @@ -359,21 +359,22 @@ public ShuffleManager(InputContext inputContext, Configuration conf, int numInpu * 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. + * A negative report is invalid and ignored. 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) { + if (numRecords < 0 || (reportedInputSet.get(inputIndex) && numRecords <= lastReported)) { return; } + + reportedInputSet.set(inputIndex); inputRecordsReported += numRecords - lastReported; numRecordsPerInput[inputIndex] = numRecords; + approximateInputRecords.setValue( extrapolateTotal(inputRecordsReported, reportedInputSet.cardinality(), numInputs)); } 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 54fe59e120..0edc21e167 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 @@ -75,6 +75,7 @@ import org.apache.tez.runtime.library.common.sort.impl.TezSpillRecord; import org.apache.tez.runtime.library.partitioner.HashPartitioner; import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads; +import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads.DataMovementEventPayloadProto; import com.google.common.collect.Lists; import com.google.protobuf.ByteString; @@ -296,9 +297,8 @@ public void testNonStartedSinglePartitionOutputReportsZeroRecords() throws Excep ShuffleUtils.generateEventsForNonStartedOutput(events, 1, outputContext, false, true, TezCommonUtils.newBestCompressionDeflater()); - ShuffleUserPayloads.DataMovementEventPayloadProto proto = - ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( - ((CompositeDataMovementEvent) events.get(0)).getUserPayload())); + DataMovementEventPayloadProto proto = DataMovementEventPayloadProto.parseFrom( + ByteString.copyFrom(((CompositeDataMovementEvent) events.get(0)).getUserPayload())); assertTrue(proto.hasNumRecord()); assertEquals(0, proto.getNumRecord()); } @@ -310,9 +310,8 @@ public void testNonStartedMultiPartitionOutputReportsNoRecordCount() throws Exce ShuffleUtils.generateEventsForNonStartedOutput(events, 10, outputContext, false, true, TezCommonUtils.newBestCompressionDeflater()); - ShuffleUserPayloads.DataMovementEventPayloadProto proto = - ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom( - ((CompositeDataMovementEvent) events.get(0)).getUserPayload())); + DataMovementEventPayloadProto proto = DataMovementEventPayloadProto.parseFrom( + ByteString.copyFrom(((CompositeDataMovementEvent) events.get(0)).getUserPayload())); assertFalse(proto.hasNumRecord()); } 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 c49b7f74d5..131eb26f48 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 @@ -453,7 +453,10 @@ public void testACountAboveIntMaxSurvivesTheWire() throws IOException { verify(shuffleManager).updateApproximateInputRecords(eq(1), eq(3_000_000_000L)); } - /** A payload with no record count -- any multi-partition writer -- reports nothing. */ + /** + * A payload that carries no record count reports nothing: the unset field is not read as a + * zero, which would count the input as having written no rows. + */ @Test @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS) public void testAPayloadWithoutARecordCountReportsNothing() 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 20d4587cfc..e0debb095f 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 @@ -245,8 +245,17 @@ public void testAReportThatLostTheRaceIsIgnored() throws Exception { assertEquals(2000L, approximateInputRecords(inputContext)); } - private long approximateInputRecords(InputContext inputContext) { - return inputContext.getCounters().findCounter(TaskCounter.APPROXIMATE_INPUT_RECORDS).getValue(); + /** + * A negative count is not a report: it neither enters the sum nor marks the input as having + * reported, so 100 from the other input is still extrapolated over both, not 95 over two. + */ + @Test + public void testANegativeReportIsIgnored() throws Exception { + InputContext inputContext = createInputContext(); + ShuffleManager shuffleManager = createShuffleManager(inputContext, 2); + shuffleManager.updateApproximateInputRecords(0, -5); + shuffleManager.updateApproximateInputRecords(1, 100); + assertEquals(200L, approximateInputRecords(inputContext)); } private InputContext createInputContext() throws IOException { @@ -397,6 +406,10 @@ private Event createDataMovementEvent(String host, int srcIndex, int targetIndex return dme; } + private long approximateInputRecords(InputContext inputContext) { + return inputContext.getCounters().findCounter(TaskCounter.APPROXIMATE_INPUT_RECORDS).getValue(); + } + private static class ShuffleManagerForTest extends ShuffleManager { public ShuffleManagerForTest(InputContext inputContext, Configuration conf, int numInputs, int bufferSize, boolean ifileReadAheadEnabled, 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 29991818c1..8fb045bae8 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 @@ -440,21 +440,21 @@ public void testTextMixedRecordsWithoutFinalMerge(boolean shouldCompress, Report @ParameterizedTest(name = "test[{0}, {1}]") @MethodSource("data") @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS) - public void testLargeRecords_SinglePartition(boolean shouldCompress, + public void testLargeRecordsSinglePartition(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. + * An output that fits one buffer spills only at close, so the events built in that same call + * are the first to see its records: the DME and the VertexManager event must both carry the + * full count. */ @ParameterizedTest(name = "test[{0}, {1}]") @MethodSource("data") @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS) - public void testFinalEventRecordCount_SinglePartition(boolean shouldCompress, + public void testFinalEventRecordCountSinglePartition(boolean shouldCompress, ReportPartitionStats reportPartitionStats) throws IOException, InterruptedException { setupInit(shouldCompress, reportPartitionStats); textTest(5, 1, 2048, 0, 0, 0, true, false); @@ -468,15 +468,18 @@ public void testFinalEventRecordCount_SinglePartition(boolean shouldCompress, * 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 { + OutputContext outputContext, List closeEvents, int numRecordsWritten) + 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() + long numRecordsCounted = kvWriter.outputRecordsCounter.getValue() + kvWriter.outputLargeRecordsCounter.getValue(); + assertEquals(numRecordsWritten, numRecordsCounted, + "outputRecordsCounter + outputLargeRecordsCounter must count every record written"); long lastDmeCount = -1; long lastVmCount = -1; for (Event event : allEvents) { @@ -494,9 +497,12 @@ private void assertLastEventsCountEveryRecord(UnorderedPartitionedKVWriter kvWri lastVmCount = vme.getNumRecord(); } } - assertTrue(lastDmeCount >= 0, "expected a DME carrying num_record"); - assertEquals(expected, lastDmeCount, "DME"); - assertEquals(expected, lastVmCount, "VertexManagerEvent"); + assertTrue(lastDmeCount >= 0, + "no DataMovementEvent carried num_record; a single-partition writer must set it"); + assertEquals(numRecordsWritten, lastDmeCount, + "num_record on the last DataMovementEvent must count every record, large ones included"); + assertEquals(numRecordsWritten, lastVmCount, + "num_record on the VertexManagerEvent must count every record, large ones included"); } public void textTest(int numRegularRecords, int numPartitions, long availableMemory, @@ -604,7 +610,7 @@ public void textTest(int numRegularRecords, int numPartitions, long availableMem verify(outputContext, never()).reportFailure(any(), any(), any()); if (numPartitions == 1) { - assertLastEventsCountEveryRecord(kvWriter, outputContext, events); + assertLastEventsCountEveryRecord(kvWriter, outputContext, events, numRecordsWritten); } if (!pipeliningEnabled) { @@ -890,6 +896,9 @@ private void baseTestWithPipelinedTransfer(int numRecords, int numPartitions, Se if (numPartitions == 1) { assertFalse(kvWriter.skipBuffers); + // With no records nothing spills and finalSpill's empty-buffer branch emits the events; + // they still carry num_record, as 0, so the consumer counts the input as having reported. + assertLastEventsCountEveryRecord(kvWriter, outputContext, lastEvents, numRecordsWritten); } //no events are sent to kvWriter upon close with pipelining