[vpj] Add HLL-based repush pipeline integrity verification

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

@@ -80,6 +80,8 @@ public class PushJobSetting implements Serializable {
   public boolean extendedSchemaValidityCheckEnabled;
   /** Refer {@link VenicePushJobConstants#COMPRESSION_METRIC_COLLECTION_ENABLED} **/
   public boolean compressionMetricCollectionEnabled;
+  public boolean repushHllVerificationEnabled;
+  public double repushHllErrorTolerance;
   public boolean repushTTLEnabled;
   public boolean isCompliancePush;
   // specify time to drop stale records.

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

@@ -28,6 +28,7 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_EXTENDED_SCHEMA_VALIDITY_CHECK_ENABLED;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_JOB_STATUS_IN_UNKNOWN_STATE_TIMEOUT_MS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_POLL_STATUS_INTERVAL_MS;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_REPUSH_HLL_ERROR_TOLERANCE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_RE_PUSH_REWIND_IN_SECONDS_OVERRIDE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_SSL_ENABLED;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFER_VERSION_SWAP;
@@ -63,6 +64,8 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.POLL_STATUS_RETRY_ATTEMPTS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_TIMEOUT_OVERRIDE_MS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_TO_SEPARATE_REALTIME_TOPIC;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_HLL_ERROR_TOLERANCE;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_HLL_VERIFICATION_ENABLED;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_ENABLE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_SECONDS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_START_TIMESTAMP;
@@ -418,6 +421,9 @@ private PushJobSetting getPushJobSetting(VeniceProperties props) {
         props.getBoolean(KAFKA_INPUT_COMPRESSION_BUILD_NEW_DICT_ENABLED, true);
     pushJobSettingToReturn.suppressEndOfPushMessage = props.getBoolean(SUPPRESS_END_OF_PUSH_MESSAGE, false);
     pushJobSettingToReturn.deferVersionSwap = props.getBoolean(DEFER_VERSION_SWAP, false);
+    pushJobSettingToReturn.repushHllVerificationEnabled = props.getBoolean(REPUSH_HLL_VERIFICATION_ENABLED, true);
+    pushJobSettingToReturn.repushHllErrorTolerance =
+        props.getDouble(REPUSH_HLL_ERROR_TOLERANCE, DEFAULT_REPUSH_HLL_ERROR_TOLERANCE);
     pushJobSettingToReturn.repushTTLEnabled = props.getBoolean(REPUSH_TTL_ENABLE, false);
     pushJobSettingToReturn.repushUseFallbackValueSchemaId =
         props.getBoolean(REPUSH_USE_FALLBACK_VALUE_SCHEMA_ID, false);
@@ -917,7 +923,8 @@ public void run() {
 
         if (!pushJobSetting.suppressEndOfPushMessage) {
           if (pushJobSetting.sendControlMessagesDirectly) {
-            getVeniceWriter(pushJobSetting).broadcastEndOfPush(Collections.emptyMap());
+            Map<Integer, Long> partitionRecordCounts = getPerPartitionRecordCounts();
+            getVeniceWriter(pushJobSetting).broadcastEndOfPush(Collections.emptyMap(), partitionRecordCounts);
           } else {
             controllerClient.writeEndOfPush(pushJobSetting.storeName, pushJobSetting.version);
           }
@@ -1742,6 +1749,17 @@ void updatePushJobDetailsWithCheckpoint(PushJobCheckpoints checkpoint) {
     pushJobDetails.pushJobLatestCheckpoint = checkpoint.getValue();
   }
 
+  private Map<Integer, Long> getPerPartitionRecordCounts() {
+    if (dataWriterComputeJob == null) {
+      return Collections.emptyMap();
+    }
+    DataWriterTaskTracker tracker = dataWriterComputeJob.getTaskTracker();
+    if (tracker == null) {
+      return Collections.emptyMap();
+    }
+    return tracker.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/mapreduce/counter/MRJobCounterHelper.java

@@ -1,5 +1,8 @@
 package com.linkedin.venice.hadoop.mapreduce.counter;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import org.apache.hadoop.mapred.Counters;
 import org.apache.hadoop.mapred.Reporter;
 
@@ -39,6 +42,8 @@ public class MRJobCounterHelper {
   private static final String COUNTER_GROUP_KAFKA_INPUT_FORMAT = "KafkaInputFormat";
   private static final String COUNTER_PUT_OR_DELETE_RECORDS = "put or delete records";
 
+  private static final String PER_PARTITION_RECORD_COUNT_GROUP = "Per Partition Record Count";
+
   private static final String REPUSH_TTL_FILTERED_COUNT = "Repush ttl filtered count";
 
   public static final GroupAndCounterNames WRITE_ACL_FAILURE_GROUP_COUNTER_NAME =
@@ -280,6 +285,32 @@ public static void incrRepushTtlFilterCount(Reporter reporter, long amount) {
     incrAmountWithGroupCounterName(reporter, REPUSH_TTL_FILTER_COUNT_GROUP_COUNTER_NAME, amount);
   }
 
+  public static void incrPartitionRecordCount(Reporter reporter, int partition, long amount) {
+    if (reporter == null || reporter.equals(Reporter.NULL) || amount == 0) {
+      return;
+    }
+    Counters.Counter counter = reporter.getCounter(PER_PARTITION_RECORD_COUNT_GROUP, String.valueOf(partition));
+    if (counter != null) {
+      counter.increment(amount);
+    }
+  }
+
+  public static Map<Integer, Long> getPerPartitionRecordCounts(Counters counters) {
+    if (counters == null) {
+      return Collections.emptyMap();
+    }
+    Map<Integer, Long> result = new HashMap<>();
+    for (Counters.Group group: counters) {
+      if (group.getName().equals(PER_PARTITION_RECORD_COUNT_GROUP)) {
+        for (Counters.Counter counter: group) {
+          result.put(Integer.parseInt(counter.getName()), counter.getValue());
+        }
+        break;
+      }
+    }
+    return result;
+  }
+
   /**
    * Bundle counter group name and counter name in this POJO
    */

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

@@ -2,6 +2,7 @@
 
 import com.linkedin.venice.hadoop.mapreduce.counter.MRJobCounterHelper;
 import com.linkedin.venice.hadoop.task.datawriter.DataWriterTaskTracker;
+import java.util.Map;
 import org.apache.hadoop.mapred.Counters;
 
 
@@ -89,4 +90,9 @@ public long getTotalPutOrDeleteRecordsCount() {
   public long getIncrementalPushThrottledTimeMs() {
     return MRJobCounterHelper.getIncrementalPushThrottleTimeMs(counters);
   }
+
+  @Override
+  public Map<Integer, Long> getPerPartitionRecordCounts() {
+    return MRJobCounterHelper.getPerPartitionRecordCounts(this.counters);
+  }
 }

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

@@ -88,6 +88,11 @@ public void trackRecordSentToPubSub() {
     MRJobCounterHelper.incrOutputRecordCount(reporter, 1);
   }
 
+  @Override
+  public void trackRecordSentToPubSubForPartition(int partition) {
+    MRJobCounterHelper.incrPartitionRecordCount(this.reporter, partition, 1);
+  }
+
   @Override
   public void trackDuplicateKeyWithDistinctValue(int count) {
     MRJobCounterHelper.incrDuplicateKeyWithDistinctValue(reporter, count);

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

@@ -498,6 +498,7 @@ private void sendMessageToKafka(
     messageSent++;
     telemetry();
     dataWriterTaskTracker.trackRecordSentToPubSub();
+    dataWriterTaskTracker.trackRecordSentToPubSubForPartition(getTaskId());
   }
 
   /**

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;
 
 
 /**
@@ -47,6 +49,9 @@ default void trackUncompressedRecordTooLargeFailure() {
   default void trackRecordSentToPubSub() {
   }
 
+  default void trackRecordSentToPubSubForPartition(int partition) {
+  }
+
   default void trackDuplicateKeyWithDistinctValue(int count) {
   }
 
@@ -132,4 +137,23 @@ default long getTotalPutOrDeleteRecordsCount() {
   default long getIncrementalPushThrottledTimeMs() {
     return 0;
   }
+
+  default Map<Integer, Long> getPerPartitionRecordCounts() {
+    return Collections.emptyMap();
+  }
+
+  default void trackReadSideUniqueKey(byte[] key) {
+  }
+
+  default void trackReadSideUniqueKeyForPartition(int partition, byte[] key) {
+  }
+
+  default long getReadSideUniqueKeyCountEstimate() {
+    return -1;
+  }
+
+  default Map<Integer, Long> getPerPartitionReadSideUniqueKeyCountEstimates() {
+    return Collections.emptyMap();
+  }
+
 }

clients/venice-push-job/src/main/java/com/linkedin/venice/jobs/DataWriterComputeJob.java

@@ -14,6 +14,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -160,6 +161,62 @@ public void validateJob() {
     } else {
       verifyTaskWithZeroValues(dataWriterTaskTracker);
     }
+
+    // Repush HLL verification: compare read-side unique key estimate with write-side output count
+    if (pushJobSetting.isSourceKafka && pushJobSetting.repushHllVerificationEnabled) {
+      long hllEstimate = dataWriterTaskTracker.getReadSideUniqueKeyCountEstimate();
+      if (hllEstimate > 0) {
+        long outputRecords = dataWriterTaskTracker.getOutputRecordsCount();
+        long ttlFiltered = dataWriterTaskTracker.getRepushTtlFilterCount();
+        // N.B.: ttlFiltered counts rows, not unique keys. This is valid because the source VT
+        // (after compaction) has one record per unique key, so row count == unique key count.
+        // The 5% default tolerance absorbs any imprecision from edge cases.
+        long expectedUniqueKeys = outputRecords + ttlFiltered;
+        double errorRate = Math.abs((double) (hllEstimate - expectedUniqueKeys)) / Math.max(hllEstimate, 1);
+
+        LOGGER.info(
+            "Repush HLL verification: HLL estimate={}, output={}, TTL filtered={}, "
+                + "expected={}, error={}, tolerance={}",
+            hllEstimate,
+            outputRecords,
+            ttlFiltered,
+            expectedUniqueKeys,
+            errorRate,
+            pushJobSetting.repushHllErrorTolerance);
+
+        if (errorRate > pushJobSetting.repushHllErrorTolerance) {
+          throw new VeniceException(
+              String.format(
+                  "Repush HLL verification failed: HLL estimate (%d) diverges from expected (%d) "
+                      + "by %.2f%%, exceeding tolerance %.2f%%",
+                  hllEstimate,
+                  expectedUniqueKeys,
+                  errorRate * 100,
+                  pushJobSetting.repushHllErrorTolerance * 100));
+        }
+
+        // Per-partition HLL check (diagnostic only, does not fail — higher variance per partition)
+        Map<Integer, Long> perPartitionHllEstimates =
+            dataWriterTaskTracker.getPerPartitionReadSideUniqueKeyCountEstimates();
+        Map<Integer, Long> perPartitionWriteCounts = dataWriterTaskTracker.getPerPartitionRecordCounts();
+        if (!perPartitionHllEstimates.isEmpty() && !perPartitionWriteCounts.isEmpty()) {
+          for (Map.Entry<Integer, Long> entry: perPartitionHllEstimates.entrySet()) {
+            int partition = entry.getKey();
+            long partHllEstimate = entry.getValue();
+            long partWriteCount = perPartitionWriteCounts.getOrDefault(partition, 0L);
+            double partErrorRate = Math.abs((double) (partHllEstimate - partWriteCount)) / Math.max(partHllEstimate, 1);
+            if (partErrorRate > pushJobSetting.repushHllErrorTolerance) {
+              LOGGER.error(
+                  "Per-partition HLL mismatch for partition {}: HLL estimate={}, write count={}, error={}",
+                  partition,
+                  partHllEstimate,
+                  partWriteCount,
+                  partErrorRate);
+            }
+          }
+        }
+      }
+    }
   }
 
   @VisibleForTesting

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

@@ -93,6 +93,8 @@
 import com.linkedin.venice.spark.datawriter.recordprocessor.SparkInputRecordProcessorFactory;
 import com.linkedin.venice.spark.datawriter.recordprocessor.SparkLogicalTimestampProcessor;
 import com.linkedin.venice.spark.datawriter.task.DataWriterAccumulators;
+import com.linkedin.venice.spark.datawriter.task.HyperLogLogAccumulator;
+import com.linkedin.venice.spark.datawriter.task.MapHyperLogLogAccumulator;
 import com.linkedin.venice.spark.datawriter.task.SparkDataWriterTaskTracker;
 import com.linkedin.venice.spark.datawriter.writer.SparkPartitionWriterFactory;
 import com.linkedin.venice.spark.input.kafka.ttl.SparkKafkaInputTTLFilter;
@@ -156,6 +158,8 @@ public abstract class AbstractDataWriterSparkJob extends DataWriterComputeJob {
   private SparkSession sparkSession;
   private DataWriterAccumulators accumulatorsForDataWriterJob;
   private SparkDataWriterTaskTracker taskTracker;
+  private HyperLogLogAccumulator readSideHllAccumulator;
+  private MapHyperLogLogAccumulator perPartitionReadSideHllAccumulator;
 
   @Override
   public void configure(VeniceProperties props, PushJobSetting pushJobSetting) {
@@ -166,7 +170,17 @@ public void configure(VeniceProperties props, PushJobSetting pushJobSetting) {
     Properties jobProps = new Properties();
     sparkSession.conf().getAll().foreach(entry -> jobProps.setProperty(entry._1, entry._2));
     accumulatorsForDataWriterJob = new DataWriterAccumulators(sparkSession);
-    taskTracker = new SparkDataWriterTaskTracker(accumulatorsForDataWriterJob);
+    if (pushJobSetting.repushHllVerificationEnabled && pushJobSetting.isSourceKafka) {
+      SparkContext sparkContext = sparkSession.sparkContext();
+      readSideHllAccumulator = new HyperLogLogAccumulator();
+      sparkContext.register(readSideHllAccumulator, "Repush Read-Side HLL Unique Key Count");
+      perPartitionReadSideHllAccumulator = new MapHyperLogLogAccumulator();
+      sparkContext.register(perPartitionReadSideHllAccumulator, "Repush Per-Partition Read-Side HLL");
+    }
+    taskTracker = new SparkDataWriterTaskTracker(
+        accumulatorsForDataWriterJob,
+        readSideHllAccumulator,
+        perPartitionReadSideHllAccumulator);
   }
 
   /**
@@ -388,6 +402,33 @@ private Dataset<Row> getInputDataFrame() {
     if (pushJobSetting.isSourceKafka) {
       Dataset<Row> rawKafkaInput = getKafkaInputDataFrame();
 
+      // Track read-side unique key cardinality via HLL for repush verification
+      if (pushJobSetting.repushHllVerificationEnabled && readSideHllAccumulator != null) {
+        final HyperLogLogAccumulator hllAcc = readSideHllAccumulator;
+        final MapHyperLogLogAccumulator perPartHllAcc = perPartitionReadSideHllAccumulator;
+        rawKafkaInput = rawKafkaInput
+            .mapPartitions((org.apache.spark.api.java.function.MapPartitionsFunction<Row, Row>) iterator -> {
+              return new java.util.Iterator<Row>() {
+                @Override
+                public boolean hasNext() {
+                  return iterator.hasNext();
+                }
+
+                @Override
+                public Row next() {
+                  Row row = iterator.next();
+                  byte[] key = row.getAs(KEY_COLUMN_NAME);
+                  if (key != null) {
+                    hllAcc.add(key);
+                    int partition = row.getAs(PARTITION_COLUMN_NAME);
+                    perPartHllAcc.add(new scala.Tuple2<>(partition, key));
+                  }
+                  return row;
+                }
+              };
+            }, org.apache.spark.sql.catalyst.encoders.RowEncoder.apply(rawKafkaInput.schema()));
+      }
+
       // Apply TTL filter first on RAW_PUBSUB_INPUT_TABLE_SCHEMA (if enabled)
       Dataset<Row> filteredInput = applyTTLFilter(rawKafkaInput);
 

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

@@ -30,6 +30,7 @@ public class DataWriterAccumulators implements Serializable {
   public final LongAccumulator repushTtlFilteredRecordCounter;
   public final LongAccumulator incrementalPushThrottleTimeCounter;
   public final LongAccumulator totalDuplicateKeyCounter;
+  public final MapLongAccumulator perPartitionRecordCounts;
 
   public DataWriterAccumulators(SparkSession session) {
     SparkContext sparkContext = session.sparkContext();
@@ -52,5 +53,6 @@ public DataWriterAccumulators(SparkSession session) {
     duplicateKeyWithIdenticalValueCounter = sparkContext.longAccumulator("Duplicate Key With Identical Value");
     duplicateKeyWithDistinctValueCounter = sparkContext.longAccumulator("Duplicate Key With Distinct Value");
     totalDuplicateKeyCounter = sparkContext.longAccumulator("Total Duplicate Keys (Compaction)");
+    perPartitionRecordCounts = new MapLongAccumulator(sparkContext, "Per Partition Record Counts");
   }
 }