TEZ-4757: APPROXIMATE_INPUT_RECORDS over-counts the rows an input will deliver

.codespellrc

@@ -18,4 +18,4 @@
 #
 
 [codespell]
-ignore-words-list = thirdparty,afterall,AfterAll
+ignore-words-list = thirdparty,afterall,AfterAll,atleast,atLeast

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<Event> 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) {

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

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.
+   * <p>
+   * 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 {

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java

@@ -861,13 +861,16 @@ public List<Event> 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) {

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 {

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<Event> 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<Event> 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";

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 {