[vpj] counting records produced per partition as part of push and adding it to the EOP message

Author: ymuppala

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java

@@ -917,10 +917,18 @@ public void run() {
         runJobWithKillDetection();
 
         if (!pushJobSetting.suppressEndOfPushMessage) {
+          Map<Integer, Long> partitionRecordCounts = getPerPartitionRecordCounts();
           if (pushJobSetting.sendControlMessagesDirectly) {
-            getVeniceWriter(pushJobSetting).broadcastEndOfPush(Collections.emptyMap());
+            getVeniceWriter(pushJobSetting).broadcastEndOfPush(Collections.emptyMap(), partitionRecordCounts);
           } else {
-            controllerClient.writeEndOfPush(pushJobSetting.storeName, pushJobSetting.version);
+            ControllerResponse eopResponse = controllerClient
+                .writeEndOfPush(pushJobSetting.storeName, pushJobSetting.version, partitionRecordCounts);
+            if (eopResponse.isError()) {
+              throw new VeniceException(
+                  "Failed to write End-of-Push for topic: "
+                      + Version.composeKafkaTopic(pushJobSetting.storeName, pushJobSetting.version) + ": "
+                      + eopResponse.getError());
+            }
           }
         }
       }
@@ -1752,6 +1760,15 @@ void updatePushJobDetailsWithCheckpoint(PushJobCheckpoints checkpoint) {
     pushJobDetails.pushJobLatestCheckpoint = checkpoint.getValue();
   }
 
+  private Map<Integer, Long> getPerPartitionRecordCounts() {
+    String topicName = Version.composeKafkaTopic(pushJobSetting.storeName, pushJobSetting.version);
+    if (dataWriterComputeJob == null || dataWriterComputeJob.getTaskTracker() == null) {
+      LOGGER.warn("Cannot retrieve per-partition record counts for topic: {}: no task tracker available", topicName);
+      return Collections.emptyMap();
+    }
+    return dataWriterComputeJob.getTaskTracker().getPerPartitionRecordCounts();
+  }
+
   private void updatePushJobDetailsWithDataWriterTracker() {
     if (dataWriterComputeJob == null) {
       LOGGER.info("No running job found. Skip updating push job details.");

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/AbstractPartitionWriter.java

@@ -267,6 +267,11 @@ public int getValueSchemaId() {
    * This doesn't need to be atomic since {@link #processValuesForKey(byte[], Iterator, DataWriterTaskTracker)} will be called sequentially.
    */
   private long messageSent = 0;
+
+  protected long getMessageSent() {
+    return messageSent;
+  }
+
   private final AtomicLong messageCompleted = new AtomicLong();
   private final AtomicLong messageErrored = new AtomicLong();
   private long timeOfLastReduceFunctionEndInNS = 0;

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/DataWriterTaskTracker.java

@@ -1,6 +1,8 @@
 package com.linkedin.venice.hadoop.task.datawriter;
 
 import com.linkedin.venice.hadoop.task.TaskTracker;
+import java.util.Collections;
+import java.util.Map;
 
 
 /**
@@ -132,4 +134,14 @@ default long getTotalPutOrDeleteRecordsCount() {
   default long getIncrementalPushThrottledTimeMs() {
     return 0;
   }
+
+  /**
+   * Returns per-partition record counts collected during the data writer job.
+   * For the Spark path, these are collected via {@code collect()} on the DAG output
+   *
+   * @return Map of partition ID to record count, or empty map if not available.
+   */
+  default Map<Integer, Long> getPerPartitionRecordCounts() {
+    return Collections.emptyMap();
+  }
 }

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/SparkConstants.java

@@ -18,6 +18,7 @@ public class SparkConstants {
 
   // Internal column names, hence begins with "_"
   public static final String PARTITION_COLUMN_NAME = "__partition__";
+  public static final String RECORD_COUNT_COLUMN_NAME = "__record_count__";
   public static final String SCHEMA_ID_COLUMN_NAME = "__schema_id__";
   public static final String RMD_VERSION_ID_COLUMN_NAME = "__replication_metadata_version_id__";
   public static final String OFFSET_COLUMN_NAME = "__offset__";
@@ -29,6 +30,10 @@ public class SparkConstants {
           new StructField(VALUE_COLUMN_NAME, BinaryType, true, Metadata.empty()),
           new StructField(RMD_COLUMN_NAME, BinaryType, true, Metadata.empty()) });
 
+  public static final StructType PARTITION_RECORD_COUNT_SCHEMA = new StructType(
+      new StructField[] { new StructField(PARTITION_COLUMN_NAME, IntegerType, false, Metadata.empty()),
+          new StructField(RECORD_COUNT_COLUMN_NAME, LongType, false, Metadata.empty()) });
+
   public static final StructType DEFAULT_SCHEMA_WITH_PARTITION = new StructType(
       new StructField[] { new StructField(KEY_COLUMN_NAME, BinaryType, false, Metadata.empty()),
           new StructField(VALUE_COLUMN_NAME, BinaryType, true, Metadata.empty()),

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/jobs/AbstractDataWriterSparkJob.java

@@ -20,6 +20,7 @@
 import static com.linkedin.venice.spark.SparkConstants.OFFSET;
 import static com.linkedin.venice.spark.SparkConstants.OFFSET_COLUMN_NAME;
 import static com.linkedin.venice.spark.SparkConstants.PARTITION_COLUMN_NAME;
+import static com.linkedin.venice.spark.SparkConstants.PARTITION_RECORD_COUNT_SCHEMA;
 import static com.linkedin.venice.spark.SparkConstants.RAW_PUBSUB_INPUT_TABLE_SCHEMA;
 import static com.linkedin.venice.spark.SparkConstants.REPLICATION_METADATA_PAYLOAD;
 import static com.linkedin.venice.spark.SparkConstants.RMD_COLUMN_NAME;
@@ -113,9 +114,11 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
@@ -844,7 +847,7 @@ public void runComputeJob() {
     validateRmdSchema(pushJobSetting);
 
     ExpressionEncoder<Row> rowEncoder = RowEncoder.apply(DEFAULT_SCHEMA);
-    ExpressionEncoder<Row> rowEncoderWithPartition = RowEncoder.apply(DEFAULT_SCHEMA_WITH_PARTITION);
+    ExpressionEncoder<Row> partitionRecordCountEncoder = RowEncoder.apply(PARTITION_RECORD_COUNT_SCHEMA);
     int numOutputPartitions = pushJobSetting.partitionCount;
 
     Properties jobProps = new Properties();
@@ -913,11 +916,30 @@ public void runComputeJob() {
         } finally {
           kafkaWriteMetrics.timeNs.add(System.nanoTime() - startNs);
         }
-      }, rowEncoderWithPartition);
-
-      // For VPJ, we don't care about the output from the DAG. ".count()" is an action that will trigger execution of
-      // the DAG to completion and will not copy all the rows to the driver to be more memory efficient.
-      dataFrame.count();
+      }, partitionRecordCountEncoder);
+
+      /*
+       * collect() returns exactly one (partitionId, recordCount) row per Spark partition. With
+       * speculative execution, if the original task and a speculative attempt both finish at the
+       * same time, Spark's TaskScheduler accepts the result from whichever one successfully
+       * communicates completion first and immediately kills the other; the duplicate's work is
+       * discarded to prevent duplicate data processing. So no two rows in the collected list will
+       * ever share the same partition id, and we can populate the map directly via put().
+       *
+       * The collected data volume is bounded by numPartitions (e.g. 10K partitions = ~160KB) so
+       * collectAsList() is safe.
+       */
+      String topicName = pushJobSetting.topic;
+      List<Row> partitionCountRows = dataFrame.collectAsList();
+      Map<Integer, Long> perPartitionRecordCounts = new HashMap<>(partitionCountRows.size());
+      for (Row row: partitionCountRows) {
+        perPartitionRecordCounts.put(row.getInt(0), row.getLong(1));
+      }
+      taskTracker.setPerPartitionRecordCounts(perPartitionRecordCounts);
+      LOGGER.info(
+          "Collected per-partition record counts for topic: {} ({} partitions)",
+          topicName,
+          perPartitionRecordCounts.size());
     } finally {
       // No matter what, always log the final accumulator values
       logAccumulatorValues();

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/task/SparkDataWriterTaskTracker.java

@@ -1,13 +1,17 @@
 package com.linkedin.venice.spark.datawriter.task;
 
 import com.linkedin.venice.hadoop.task.datawriter.DataWriterTaskTracker;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 
 /**
  * This class is used to track the metrics for the Spark Data Writer task.
  */
 public class SparkDataWriterTaskTracker implements DataWriterTaskTracker {
   private final DataWriterAccumulators accumulators;
+  private Map<Integer, Long> perPartitionRecordCounts = Collections.emptyMap();
 
   public SparkDataWriterTaskTracker(DataWriterAccumulators accumulators) {
     this.accumulators = accumulators;
@@ -172,4 +176,18 @@ public long getRepushTtlFilterCount() {
   public long getIncrementalPushThrottledTimeMs() {
     return accumulators.incrementalPushThrottleTimeCounter.value();
   }
+
+  /**
+   * Sets the per-partition record counts collected from the Spark DAG output via {@code collect()}.
+   */
+  public void setPerPartitionRecordCounts(Map<Integer, Long> counts) {
+    this.perPartitionRecordCounts = (counts == null || counts.isEmpty())
+        ? Collections.emptyMap()
+        : Collections.unmodifiableMap(new HashMap<>(counts));
+  }
+
+  @Override
+  public Map<Integer, Long> getPerPartitionRecordCounts() {
+    return perPartitionRecordCounts;
+  }
 }

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/writer/SparkPartitionWriter.java

@@ -85,4 +85,11 @@ void processRows(Iterator<Row> rows) {
       super.processValuesForKey(key, valueRecordsForKey.iterator(), dataWriterTaskTracker);
     }
   }
+
+  /**
+   * @return The number of records sent to PubSub by this partition writer.
+   */
+  long getRecordCount() {
+    return getMessageSent();
+  }
 }

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/writer/SparkPartitionWriterFactory.java

@@ -1,11 +1,14 @@
 package com.linkedin.venice.spark.datawriter.writer;
 
 import com.linkedin.venice.spark.datawriter.task.DataWriterAccumulators;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.Properties;
+import org.apache.spark.TaskContext;
 import org.apache.spark.api.java.function.MapPartitionsFunction;
 import org.apache.spark.broadcast.Broadcast;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
 
 
 public class SparkPartitionWriterFactory implements MapPartitionsFunction<Row, Row> {
@@ -20,9 +23,12 @@ public SparkPartitionWriterFactory(Broadcast<Properties> jobProps, DataWriterAcc
 
   @Override
   public Iterator<Row> call(Iterator<Row> rows) throws Exception {
+    long recordCount;
     try (SparkPartitionWriter partitionWriter = new SparkPartitionWriter(jobProps.getValue(), accumulators)) {
       partitionWriter.processRows(rows);
+      recordCount = partitionWriter.getRecordCount();
     }
-    return rows;
+    int partitionId = TaskContext.get().partitionId();
+    return Collections.singletonList(RowFactory.create(partitionId, recordCount)).iterator();
   }
 }

clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/TestVenicePushJobCheckpoints.java

@@ -746,6 +746,8 @@ private void configureControllerClientMock(
     ControllerResponse controllerResponse = mock(ControllerResponse.class);
     when(controllerResponse.isError()).thenReturn(false);
     when(controllerClient.sendPushJobDetails(anyString(), anyInt(), any(byte[].class))).thenReturn(controllerResponse);
+    when(controllerClient.writeEndOfPush(anyString(), anyInt())).thenReturn(controllerResponse);
+    when(controllerClient.writeEndOfPush(anyString(), anyInt(), any())).thenReturn(controllerResponse);
   }
 
   private void configureClusterDiscoverControllerClient(ControllerClient controllerClient) {

clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobTest.java

@@ -809,6 +809,7 @@ private ControllerClient getClient(Consumer<StoreInfo> storeInfo, boolean applyF
 
     ControllerResponse response = new ControllerResponse();
     doReturn(response).when(client).writeEndOfPush(anyString(), anyInt());
+    doReturn(response).when(client).writeEndOfPush(anyString(), anyInt(), any());
     doReturn(response).when(client).sendPushJobDetails(anyString(), anyInt(), any(byte[].class));
     return client;
   }