Skip to content

Commit 69c0c85

Browse files
sushantmanesushantmane
committed
[vpj] Add HLL-based repush pipeline integrity verification
During KIF repush, track read-side unique key cardinality via HLL and compare with write-side output count at the driver. Two-tier verification: aggregate HLL (fails push) + per-partition HLL (diagnostic log).
1 parent 3ea6a3b commit 69c0c85

10 files changed

Lines changed: 414 additions & 1 deletion

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ public class PushJobSetting implements Serializable {
8080
public boolean extendedSchemaValidityCheckEnabled;
8181
/** Refer {@link VenicePushJobConstants#COMPRESSION_METRIC_COLLECTION_ENABLED} **/
8282
public boolean compressionMetricCollectionEnabled;
83+
public boolean repushHllVerificationEnabled;
84+
public double repushHllErrorTolerance;
8385
public boolean repushTTLEnabled;
8486
public boolean isCompliancePush;
8587
// specify time to drop stale records.

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_EXTENDED_SCHEMA_VALIDITY_CHECK_ENABLED;
2929
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_JOB_STATUS_IN_UNKNOWN_STATE_TIMEOUT_MS;
3030
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_POLL_STATUS_INTERVAL_MS;
31+
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_REPUSH_HLL_ERROR_TOLERANCE;
3132
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_RE_PUSH_REWIND_IN_SECONDS_OVERRIDE;
3233
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_SSL_ENABLED;
3334
import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFER_VERSION_SWAP;
@@ -63,6 +64,8 @@
6364
import static com.linkedin.venice.vpj.VenicePushJobConstants.POLL_STATUS_RETRY_ATTEMPTS;
6465
import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_TIMEOUT_OVERRIDE_MS;
6566
import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_TO_SEPARATE_REALTIME_TOPIC;
67+
import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_HLL_ERROR_TOLERANCE;
68+
import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_HLL_VERIFICATION_ENABLED;
6669
import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_ENABLE;
6770
import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_SECONDS;
6871
import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_START_TIMESTAMP;
@@ -418,6 +421,9 @@ private PushJobSetting getPushJobSetting(VeniceProperties props) {
418421
props.getBoolean(KAFKA_INPUT_COMPRESSION_BUILD_NEW_DICT_ENABLED, true);
419422
pushJobSettingToReturn.suppressEndOfPushMessage = props.getBoolean(SUPPRESS_END_OF_PUSH_MESSAGE, false);
420423
pushJobSettingToReturn.deferVersionSwap = props.getBoolean(DEFER_VERSION_SWAP, false);
424+
pushJobSettingToReturn.repushHllVerificationEnabled = props.getBoolean(REPUSH_HLL_VERIFICATION_ENABLED, true);
425+
pushJobSettingToReturn.repushHllErrorTolerance =
426+
props.getDouble(REPUSH_HLL_ERROR_TOLERANCE, DEFAULT_REPUSH_HLL_ERROR_TOLERANCE);
421427
pushJobSettingToReturn.repushTTLEnabled = props.getBoolean(REPUSH_TTL_ENABLE, false);
422428
pushJobSettingToReturn.repushUseFallbackValueSchemaId =
423429
props.getBoolean(REPUSH_USE_FALLBACK_VALUE_SCHEMA_ID, false);

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,4 +141,19 @@ default long getIncrementalPushThrottledTimeMs() {
141141
default Map<Integer, Long> getPerPartitionRecordCounts() {
142142
return Collections.emptyMap();
143143
}
144+
145+
default void trackReadSideUniqueKey(byte[] key) {
146+
}
147+
148+
default void trackReadSideUniqueKeyForPartition(int partition, byte[] key) {
149+
}
150+
151+
default long getReadSideUniqueKeyCountEstimate() {
152+
return -1;
153+
}
154+
155+
default Map<Integer, Long> getPerPartitionReadSideUniqueKeyCountEstimates() {
156+
return Collections.emptyMap();
157+
}
158+
144159
}

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

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import java.util.Arrays;
1515
import java.util.Collections;
1616
import java.util.List;
17+
import java.util.Map;
1718
import org.apache.logging.log4j.LogManager;
1819
import org.apache.logging.log4j.Logger;
1920

@@ -160,6 +161,62 @@ public void validateJob() {
160161
} else {
161162
verifyTaskWithZeroValues(dataWriterTaskTracker);
162163
}
164+
165+
// Repush HLL verification: compare read-side unique key estimate with write-side output count
166+
if (pushJobSetting.isSourceKafka && pushJobSetting.repushHllVerificationEnabled) {
167+
long hllEstimate = dataWriterTaskTracker.getReadSideUniqueKeyCountEstimate();
168+
if (hllEstimate > 0) {
169+
long outputRecords = dataWriterTaskTracker.getOutputRecordsCount();
170+
long ttlFiltered = dataWriterTaskTracker.getRepushTtlFilterCount();
171+
// N.B.: ttlFiltered counts rows, not unique keys. This is valid because the source VT
172+
// (after compaction) has one record per unique key, so row count == unique key count.
173+
// The 5% default tolerance absorbs any imprecision from edge cases.
174+
long expectedUniqueKeys = outputRecords + ttlFiltered;
175+
double errorRate = Math.abs((double) (hllEstimate - expectedUniqueKeys)) / Math.max(hllEstimate, 1);
176+
177+
LOGGER.info(
178+
"Repush HLL verification: HLL estimate={}, output={}, TTL filtered={}, "
179+
+ "expected={}, error={}, tolerance={}",
180+
hllEstimate,
181+
outputRecords,
182+
ttlFiltered,
183+
expectedUniqueKeys,
184+
errorRate,
185+
pushJobSetting.repushHllErrorTolerance);
186+
187+
if (errorRate > pushJobSetting.repushHllErrorTolerance) {
188+
throw new VeniceException(
189+
String.format(
190+
"Repush HLL verification failed: HLL estimate (%d) diverges from expected (%d) "
191+
+ "by %.2f%%, exceeding tolerance %.2f%%",
192+
hllEstimate,
193+
expectedUniqueKeys,
194+
errorRate * 100,
195+
pushJobSetting.repushHllErrorTolerance * 100));
196+
}
197+
198+
// Per-partition HLL check (diagnostic only, does not fail — higher variance per partition)
199+
Map<Integer, Long> perPartitionHllEstimates =
200+
dataWriterTaskTracker.getPerPartitionReadSideUniqueKeyCountEstimates();
201+
Map<Integer, Long> perPartitionWriteCounts = dataWriterTaskTracker.getPerPartitionRecordCounts();
202+
if (!perPartitionHllEstimates.isEmpty() && !perPartitionWriteCounts.isEmpty()) {
203+
for (Map.Entry<Integer, Long> entry: perPartitionHllEstimates.entrySet()) {
204+
int partition = entry.getKey();
205+
long partHllEstimate = entry.getValue();
206+
long partWriteCount = perPartitionWriteCounts.getOrDefault(partition, 0L);
207+
double partErrorRate = Math.abs((double) (partHllEstimate - partWriteCount)) / Math.max(partHllEstimate, 1);
208+
if (partErrorRate > pushJobSetting.repushHllErrorTolerance) {
209+
LOGGER.error(
210+
"Per-partition HLL mismatch for partition {}: HLL estimate={}, write count={}, error={}",
211+
partition,
212+
partHllEstimate,
213+
partWriteCount,
214+
partErrorRate);
215+
}
216+
}
217+
}
218+
}
219+
}
163220
}
164221

165222
@VisibleForTesting

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

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,8 @@
9393
import com.linkedin.venice.spark.datawriter.recordprocessor.SparkInputRecordProcessorFactory;
9494
import com.linkedin.venice.spark.datawriter.recordprocessor.SparkLogicalTimestampProcessor;
9595
import com.linkedin.venice.spark.datawriter.task.DataWriterAccumulators;
96+
import com.linkedin.venice.spark.datawriter.task.HyperLogLogAccumulator;
97+
import com.linkedin.venice.spark.datawriter.task.MapHyperLogLogAccumulator;
9698
import com.linkedin.venice.spark.datawriter.task.SparkDataWriterTaskTracker;
9799
import com.linkedin.venice.spark.datawriter.writer.SparkPartitionWriterFactory;
98100
import com.linkedin.venice.spark.input.kafka.ttl.SparkKafkaInputTTLFilter;
@@ -156,6 +158,8 @@ public abstract class AbstractDataWriterSparkJob extends DataWriterComputeJob {
156158
private SparkSession sparkSession;
157159
private DataWriterAccumulators accumulatorsForDataWriterJob;
158160
private SparkDataWriterTaskTracker taskTracker;
161+
private HyperLogLogAccumulator readSideHllAccumulator;
162+
private MapHyperLogLogAccumulator perPartitionReadSideHllAccumulator;
159163

160164
@Override
161165
public void configure(VeniceProperties props, PushJobSetting pushJobSetting) {
@@ -166,7 +170,17 @@ public void configure(VeniceProperties props, PushJobSetting pushJobSetting) {
166170
Properties jobProps = new Properties();
167171
sparkSession.conf().getAll().foreach(entry -> jobProps.setProperty(entry._1, entry._2));
168172
accumulatorsForDataWriterJob = new DataWriterAccumulators(sparkSession);
169-
taskTracker = new SparkDataWriterTaskTracker(accumulatorsForDataWriterJob);
173+
if (pushJobSetting.repushHllVerificationEnabled && pushJobSetting.isSourceKafka) {
174+
SparkContext sparkContext = sparkSession.sparkContext();
175+
readSideHllAccumulator = new HyperLogLogAccumulator();
176+
sparkContext.register(readSideHllAccumulator, "Repush Read-Side HLL Unique Key Count");
177+
perPartitionReadSideHllAccumulator = new MapHyperLogLogAccumulator();
178+
sparkContext.register(perPartitionReadSideHllAccumulator, "Repush Per-Partition Read-Side HLL");
179+
}
180+
taskTracker = new SparkDataWriterTaskTracker(
181+
accumulatorsForDataWriterJob,
182+
readSideHllAccumulator,
183+
perPartitionReadSideHllAccumulator);
170184
}
171185

172186
/**
@@ -388,6 +402,33 @@ private Dataset<Row> getInputDataFrame() {
388402
if (pushJobSetting.isSourceKafka) {
389403
Dataset<Row> rawKafkaInput = getKafkaInputDataFrame();
390404

405+
// Track read-side unique key cardinality via HLL for repush verification
406+
if (pushJobSetting.repushHllVerificationEnabled && readSideHllAccumulator != null) {
407+
final HyperLogLogAccumulator hllAcc = readSideHllAccumulator;
408+
final MapHyperLogLogAccumulator perPartHllAcc = perPartitionReadSideHllAccumulator;
409+
rawKafkaInput = rawKafkaInput
410+
.mapPartitions((org.apache.spark.api.java.function.MapPartitionsFunction<Row, Row>) iterator -> {
411+
return new java.util.Iterator<Row>() {
412+
@Override
413+
public boolean hasNext() {
414+
return iterator.hasNext();
415+
}
416+
417+
@Override
418+
public Row next() {
419+
Row row = iterator.next();
420+
byte[] key = row.getAs(KEY_COLUMN_NAME);
421+
if (key != null) {
422+
hllAcc.add(key);
423+
int partition = row.getAs(PARTITION_COLUMN_NAME);
424+
perPartHllAcc.add(new scala.Tuple2<>(partition, key));
425+
}
426+
return row;
427+
}
428+
};
429+
}, org.apache.spark.sql.catalyst.encoders.RowEncoder.apply(rawKafkaInput.schema()));
430+
}
431+
391432
// Apply TTL filter first on RAW_PUBSUB_INPUT_TABLE_SCHEMA (if enabled)
392433
Dataset<Row> filteredInput = applyTTLFilter(rawKafkaInput);
393434

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package com.linkedin.venice.spark.datawriter.task;
2+
3+
import com.linkedin.venice.utils.HyperLogLogSketch;
4+
import org.apache.spark.util.AccumulatorV2;
5+
6+
7+
/**
8+
* A Spark {@link AccumulatorV2} that wraps {@link HyperLogLogSketch} for distributed
9+
* cardinality estimation across Spark tasks.
10+
*
11+
* Each task gets a fresh copy via {@link #copy()}, adds keys during processing,
12+
* and the driver merges all task-level sketches via {@link #merge(AccumulatorV2)}.
13+
*/
14+
public class HyperLogLogAccumulator extends AccumulatorV2<byte[], Long> {
15+
private static final long serialVersionUID = 1L;
16+
17+
private HyperLogLogSketch sketch;
18+
19+
public HyperLogLogAccumulator() {
20+
this.sketch = new HyperLogLogSketch();
21+
}
22+
23+
@Override
24+
public boolean isZero() {
25+
return sketch.isEmpty();
26+
}
27+
28+
@Override
29+
public AccumulatorV2<byte[], Long> copy() {
30+
HyperLogLogAccumulator newAcc = new HyperLogLogAccumulator();
31+
newAcc.sketch = this.sketch.copy();
32+
return newAcc;
33+
}
34+
35+
@Override
36+
public void reset() {
37+
sketch.reset();
38+
}
39+
40+
@Override
41+
public void add(byte[] key) {
42+
sketch.add(key);
43+
}
44+
45+
@Override
46+
public void merge(AccumulatorV2<byte[], Long> other) {
47+
sketch.merge(((HyperLogLogAccumulator) other).sketch);
48+
}
49+
50+
@Override
51+
public Long value() {
52+
return sketch.estimate();
53+
}
54+
}

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package com.linkedin.venice.spark.datawriter.task;
2+
3+
import java.util.Collections;
4+
import java.util.Map;
5+
import java.util.concurrent.ConcurrentHashMap;
6+
import org.apache.spark.util.AccumulatorV2;
7+
import scala.Tuple2;
8+
9+
10+
/**
11+
* A Spark accumulator that maintains per-partition HyperLogLog sketches for estimating
12+
* the unique key cardinality within each partition independently.
13+
*
14+
* Used during repush when source and destination partition counts match, enabling
15+
* per-partition verification that catches partition-level record loss which aggregate
16+
* HLL comparison would mask.
17+
*
18+
* Each partition gets its own HLL sketch (16KB). For 1000 partitions, total memory is ~16MB per task.
19+
*/
20+
public class MapHyperLogLogAccumulator extends AccumulatorV2<Tuple2<Integer, byte[]>, Map<Integer, Long>> {
21+
private static final long serialVersionUID = 1L;
22+
23+
private final ConcurrentHashMap<Integer, HyperLogLogAccumulator> hllMap = new ConcurrentHashMap<>();
24+
25+
@Override
26+
public boolean isZero() {
27+
return hllMap.isEmpty();
28+
}
29+
30+
@Override
31+
public AccumulatorV2<Tuple2<Integer, byte[]>, Map<Integer, Long>> copy() {
32+
MapHyperLogLogAccumulator newAcc = new MapHyperLogLogAccumulator();
33+
hllMap.forEach((partition, hll) -> {
34+
newAcc.hllMap.put(partition, (HyperLogLogAccumulator) hll.copy());
35+
});
36+
return newAcc;
37+
}
38+
39+
@Override
40+
public void reset() {
41+
hllMap.clear();
42+
}
43+
44+
@Override
45+
public void add(Tuple2<Integer, byte[]> v) {
46+
hllMap.computeIfAbsent(v._1(), k -> new HyperLogLogAccumulator()).add(v._2());
47+
}
48+
49+
@Override
50+
public void merge(AccumulatorV2<Tuple2<Integer, byte[]>, Map<Integer, Long>> other) {
51+
MapHyperLogLogAccumulator otherAcc = (MapHyperLogLogAccumulator) other;
52+
otherAcc.hllMap.forEach((partition, otherHll) -> {
53+
hllMap.merge(partition, otherHll, (existing, incoming) -> {
54+
existing.merge(incoming);
55+
return existing;
56+
});
57+
});
58+
}
59+
60+
@Override
61+
public Map<Integer, Long> value() {
62+
ConcurrentHashMap<Integer, Long> result = new ConcurrentHashMap<>();
63+
hllMap.forEach((partition, hll) -> result.put(partition, hll.value()));
64+
return Collections.unmodifiableMap(result);
65+
}
66+
}

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.linkedin.venice.spark.datawriter.task;
22

33
import com.linkedin.venice.hadoop.task.datawriter.DataWriterTaskTracker;
4+
import java.util.Collections;
45
import java.util.Map;
56

67

@@ -9,9 +10,20 @@
910
*/
1011
public class SparkDataWriterTaskTracker implements DataWriterTaskTracker {
1112
private final DataWriterAccumulators accumulators;
13+
private final HyperLogLogAccumulator readSideHllAccumulator;
14+
private final MapHyperLogLogAccumulator perPartitionReadSideHllAccumulator;
1215

1316
public SparkDataWriterTaskTracker(DataWriterAccumulators accumulators) {
17+
this(accumulators, null, null);
18+
}
19+
20+
public SparkDataWriterTaskTracker(
21+
DataWriterAccumulators accumulators,
22+
HyperLogLogAccumulator readSideHllAccumulator,
23+
MapHyperLogLogAccumulator perPartitionReadSideHllAccumulator) {
1424
this.accumulators = accumulators;
25+
this.readSideHllAccumulator = readSideHllAccumulator;
26+
this.perPartitionReadSideHllAccumulator = perPartitionReadSideHllAccumulator;
1527
}
1628

1729
@Override
@@ -183,4 +195,31 @@ public long getIncrementalPushThrottledTimeMs() {
183195
public Map<Integer, Long> getPerPartitionRecordCounts() {
184196
return accumulators.perPartitionRecordCounts.value();
185197
}
198+
199+
@Override
200+
public void trackReadSideUniqueKey(byte[] key) {
201+
if (readSideHllAccumulator != null) {
202+
readSideHllAccumulator.add(key);
203+
}
204+
}
205+
206+
@Override
207+
public void trackReadSideUniqueKeyForPartition(int partition, byte[] key) {
208+
if (perPartitionReadSideHllAccumulator != null) {
209+
perPartitionReadSideHllAccumulator.add(new scala.Tuple2<>(partition, key));
210+
}
211+
}
212+
213+
@Override
214+
public long getReadSideUniqueKeyCountEstimate() {
215+
return readSideHllAccumulator != null ? readSideHllAccumulator.value() : -1;
216+
}
217+
218+
@Override
219+
public Map<Integer, Long> getPerPartitionReadSideUniqueKeyCountEstimates() {
220+
return perPartitionReadSideHllAccumulator != null
221+
? perPartitionReadSideHllAccumulator.value()
222+
: Collections.emptyMap();
223+
}
224+
186225
}

‎clients/venice-push-job/src/main/java/com/linkedin/venice/vpj/VenicePushJobConstants.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,14 @@ private VenicePushJobConstants() {
406406
*/
407407
public static final String REPUSH_USE_FALLBACK_VALUE_SCHEMA_ID = "repush.use.fallback.value.schema.id";
408408

409+
/**
410+
* When enabled, repush jobs use HyperLogLog to estimate the unique key cardinality of the source
411+
* version topic and compare it against the write-side output count for pipeline integrity verification.
412+
*/
413+
public static final String REPUSH_HLL_VERIFICATION_ENABLED = "repush.hll.verification.enabled";
414+
public static final String REPUSH_HLL_ERROR_TOLERANCE = "repush.hll.error.tolerance";
415+
public static final double DEFAULT_REPUSH_HLL_ERROR_TOLERANCE = 0.05;
416+
409417
public static final String REPUSH_TTL_ENABLE = "repush.ttl.enable";
410418
public static final String REPUSH_TTL_POLICY = "repush.ttl.policy";
411419
public static final String REPUSH_TTL_SECONDS = "repush.ttl.seconds";

0 commit comments

Comments
 (0)