Addressed comments and fixed re-push and compression support

Author: xunyin8

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java

@@ -549,7 +549,8 @@ private PubSubMessageProcessedResult processActiveActiveMessage(
               rmdWithValueSchemaID,
               valueManifestContainer,
               null,
-              null));
+              null,
+              (schemaId) -> storeDeserializerCache.getDeserializer(schemaId, schemaId)));
     } else {
       validatePostOperationResultsAndRecord(mergeConflictResult, offsetSumPreOperation, recordTimestampsPreOperation);
 
@@ -589,7 +590,8 @@ private PubSubMessageProcessedResult processActiveActiveMessage(
               rmdWithValueSchemaID,
               valueManifestContainer,
               updatedValueBytes,
-              updatedRmdBytes));
+              updatedRmdBytes,
+              (schemaId) -> storeDeserializerCache.getDeserializer(schemaId, schemaId)));
     }
   }
 
@@ -669,8 +671,7 @@ protected void processMessageAndMaybeProduceToKafka(
         ByteBuffer oldValueBB = mergeConflictResultWrapper.getOldValueByteBufferProvider().get();
         int oldValueSchemaId =
             oldValueBB == null ? -1 : mergeConflictResultWrapper.getOldValueProvider().get().writerSchemaId();
-        Lazy<GenericRecord> newValueProvider = mergeConflictResultWrapper
-            .getNewValueProvider((schemaId) -> storeDeserializerCache.getDeserializer(schemaId, schemaId));
+        Lazy<GenericRecord> newValueProvider = mergeConflictResultWrapper.getNewValueProvider();
         queueUpVersionTopicWritesWithViewWriters(
             partitionConsumptionState,
             (viewWriter) -> viewWriter.processRecord(

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java

@@ -71,6 +71,7 @@
 import com.linkedin.venice.schema.SchemaEntry;
 import com.linkedin.venice.schema.writecompute.DerivedSchemaEntry;
 import com.linkedin.venice.serialization.AvroStoreDeserializerCache;
+import com.linkedin.venice.serializer.RecordDeserializer;
 import com.linkedin.venice.stats.StatsErrorCode;
 import com.linkedin.venice.storage.protocol.ChunkedValueManifest;
 import com.linkedin.venice.utils.ByteUtils;
@@ -3202,8 +3203,20 @@ private PubSubMessageProcessedResult processMessage(
         Put put = (Put) kafkaValue.payloadUnion;
         // Value provider should use un-compressed data.
         final ByteBuffer rawPutValue = put.putValue;
-        valueProvider =
-            Lazy.of(() -> storeDeserializerCache.getDeserializer(put.schemaId, put.schemaId).deserialize(rawPutValue));
+        final boolean needToDecompress = !partitionConsumptionState.isEndOfPushReceived();
+        valueProvider = Lazy.of(() -> {
+          RecordDeserializer<GenericRecord> recordDeserializer =
+              storeDeserializerCache.getDeserializer(put.schemaId, put.schemaId);
+          if (needToDecompress) {
+            try {
+              return recordDeserializer.deserialize(compressor.get().decompress(rawPutValue));
+            } catch (IOException e) {
+              throw new VeniceException("Unable to provide value due to decompression failure", e);
+            }
+          } else {
+            return recordDeserializer.deserialize(rawPutValue);
+          }
+        });
         put.putValue = maybeCompressData(
             consumerRecord.getTopicPartition().getPartitionNumber(),
             put.putValue,
@@ -3271,21 +3284,20 @@ private PubSubMessageProcessedResult processMessage(
 
         final byte[] updatedValueBytes;
         final ChunkedValueManifest oldValueManifest = valueManifestContainer.getManifest();
-        GenericRecord updatedValue;
+        WriteComputeResult writeComputeResult;
         try {
           long writeComputeStartTimeInNS = System.nanoTime();
 
           // Leader nodes are the only ones which process UPDATES, so it's valid to always compress and not call
           // 'maybeCompress'.
-          updatedValue = storeWriteComputeHandler.applyWriteCompute(
+          writeComputeResult = storeWriteComputeHandler.applyWriteCompute(
               currValue,
               update.schemaId,
               readerValueSchemaId,
               update.updateValue,
               update.updateSchemaId,
               readerUpdateProtocolVersion);
-          updatedValueBytes = compressor.get()
-              .compress(storeWriteComputeHandler.serializeUpdatedValue(updatedValue, readerValueSchemaId));
+          updatedValueBytes = writeComputeResult.getUpdatedValueBytes();
           hostLevelIngestionStats
               .recordWriteComputeUpdateLatency(LatencyUtils.getElapsedTimeFromNSToMS(writeComputeStartTimeInNS));
         } catch (Exception e) {
@@ -3323,7 +3335,11 @@ private PubSubMessageProcessedResult processMessage(
           updatedPut.putValue = updateValueWithSchemaId;
           updatedPut.schemaId = readerValueSchemaId;
           return new PubSubMessageProcessedResult(
-              new WriteComputeResultWrapper(updatedPut, oldValueManifest, false, Lazy.of(() -> updatedValue)));
+              new WriteComputeResultWrapper(
+                  updatedPut,
+                  oldValueManifest,
+                  false,
+                  Lazy.of(writeComputeResult::getUpdatedValue)));
         }
       case DELETE:
         /**

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/MergeConflictResultWrapper.java

@@ -20,24 +20,53 @@ public class MergeConflictResultWrapper {
   private final Lazy<ByteBuffer> oldValueByteBufferProvider;
   private final RmdWithValueSchemaId oldRmdWithValueSchemaId;
   private final ChunkedValueManifestContainer oldValueManifestContainer;
+
+  // Serialized and potentially compressed updated value bytes
   private final ByteBuffer updatedValueBytes;
   private final ByteBuffer updatedRmdBytes;
 
+  /**
+   * Best-effort deserialized value provider that provides the updated value for PUT/UPDATE and the old value for
+   * DELETE.
+   */
+  private final Lazy<GenericRecord> valueProvider;
+
   public MergeConflictResultWrapper(
       MergeConflictResult mergeConflictResult,
       Lazy<ByteBufferValueRecord<ByteBuffer>> oldValueProvider,
       Lazy<ByteBuffer> oldValueByteBufferProvider,
       RmdWithValueSchemaId oldRmdWithValueSchemaId,
       ChunkedValueManifestContainer oldValueManifestContainer,
       ByteBuffer updatedValueBytes,
-      ByteBuffer updatedRmdBytes) {
+      ByteBuffer updatedRmdBytes,
+      Function<Integer, RecordDeserializer<GenericRecord>> deserializerProvider) {
     this.mergeConflictResult = mergeConflictResult;
     this.oldValueProvider = oldValueProvider;
     this.oldValueByteBufferProvider = oldValueByteBufferProvider;
     this.oldRmdWithValueSchemaId = oldRmdWithValueSchemaId;
     this.oldValueManifestContainer = oldValueManifestContainer;
     this.updatedValueBytes = updatedValueBytes;
     this.updatedRmdBytes = updatedRmdBytes;
+    if (updatedValueBytes == null) {
+      // this is a DELETE
+      ByteBufferValueRecord<ByteBuffer> oldValue = oldValueProvider.get();
+      if (oldValue == null || oldValue.value() == null) {
+        this.valueProvider = Lazy.of(() -> null);
+      } else {
+        this.valueProvider =
+            Lazy.of(() -> deserializerProvider.apply(oldValue.writerSchemaId()).deserialize(oldValue.value()));
+      }
+    } else {
+      // this is a PUT or UPDATE
+      if (mergeConflictResult.getDeserializedValue().isPresent()) {
+        this.valueProvider = Lazy.of(() -> mergeConflictResult.getDeserializedValue().get());
+      } else {
+        // Use mergeConflictResult.getNewValue() here since updatedValueBytes could be compressed.
+        this.valueProvider = Lazy.of(
+            () -> deserializerProvider.apply(mergeConflictResult.getValueSchemaId())
+                .deserialize(mergeConflictResult.getNewValue()));
+      }
+    }
   }
 
   public MergeConflictResult getMergeConflictResult() {
@@ -74,24 +103,7 @@ public ByteBuffer getUpdatedRmdBytes() {
    *   2. returns the old value for DELETE (null for non-existent key).
    *   3. returns null if the value is not available.
    */
-  public Lazy<GenericRecord> getNewValueProvider(
-      Function<Integer, RecordDeserializer<GenericRecord>> deserializerProvider) {
-    if (updatedValueBytes == null) {
-      // this is a DELETE
-      ByteBufferValueRecord<ByteBuffer> oldValue = oldValueProvider.get();
-      if (oldValue == null || oldValue.value() == null) {
-        return Lazy.of(() -> null);
-      }
-      return Lazy.of(() -> deserializerProvider.apply(oldValue.writerSchemaId()).deserialize(oldValue.value()));
-    } else {
-      // this is a PUT or UPDATE
-      if (mergeConflictResult.getValueDeserialized().isPresent()) {
-        return Lazy.of(() -> mergeConflictResult.getValueDeserialized().get());
-      }
-      // Use mergeConflictResult.getNewValue() here and not updatedValueBytes for non-compressed value bytes.
-      return Lazy.of(
-          () -> deserializerProvider.apply(mergeConflictResult.getValueSchemaId())
-              .deserialize(mergeConflictResult.getNewValue()));
-    }
+  public Lazy<GenericRecord> getNewValueProvider() {
+    return valueProvider;
   }
 }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreWriteComputeProcessor.java

@@ -91,9 +91,9 @@ public StoreWriteComputeProcessor(
    * @param writerUpdateProtocolVersion Update protocol version used to serialize Update payload bytes.
    * @param readerUpdateProtocolVersion Update protocol version used to deserialize Update payload bytes.
    *
-   * @return Bytes of partially updated original value.
+   * @return {@link WriteComputeResult} of partially updated original value.
    */
-  public GenericRecord applyWriteCompute(
+  public WriteComputeResult applyWriteCompute(
       GenericRecord currValue,
       int writerValueSchemaId,
       int readerValueSchemaId,
@@ -110,11 +110,10 @@ public GenericRecord applyWriteCompute(
         writeComputeProcessor.updateRecord(readerSchemaContainer.getValueSchema(), currValue, writeComputeRecord);
 
     // If write compute is enabled and the record is deleted, the updatedValue will be null.
-    return updatedValue;
-  }
-
-  public byte[] serializeUpdatedValue(GenericRecord updatedValue, int readerValueSchemaId) {
-    return getValueSerializer(readerValueSchemaId).serialize(updatedValue);
+    if (updatedValue == null) {
+      return new WriteComputeResult(null, null);
+    }
+    return new WriteComputeResult(getValueSerializer(readerValueSchemaId).serialize(updatedValue), updatedValue);
   }
 
   private SchemaAndUniqueId getSchemaAndUniqueId(int valueSchemaId, int writeComputeSchemaId) {

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/WriteComputeResult.java

@@ -0,0 +1,29 @@
+package com.linkedin.davinci.kafka.consumer;
+
+import javax.annotation.Nullable;
+import org.apache.avro.generic.GenericRecord;
+
+
+/**
+ * Write compute result wrapper class holding the deserialized updated value and the serialized and potentially
+ * compressed updated value bytes.
+ */
+public class WriteComputeResult {
+  private final byte[] updatedValueBytes;
+  private final GenericRecord updatedValue;
+
+  public WriteComputeResult(byte[] updatedValueBytes, GenericRecord updatedValue) {
+    this.updatedValueBytes = updatedValueBytes;
+    this.updatedValue = updatedValue;
+  }
+
+  @Nullable
+  public byte[] getUpdatedValueBytes() {
+    return updatedValueBytes;
+  }
+
+  @Nullable
+  public GenericRecord getUpdatedValue() {
+    return updatedValue;
+  }
+}

clients/da-vinci-client/src/main/java/com/linkedin/davinci/replication/merge/MergeConflictResult.java

@@ -13,7 +13,7 @@ public class MergeConflictResult {
   private static final MergeConflictResult IGNORED_RESULT = new MergeConflictResult();
 
   private ByteBuffer newValue;
-  private Optional<GenericRecord> valueDeserialized;
+  private Optional<GenericRecord> deserializedValue;
   private int valueSchemaId;
   private final boolean updateIgnored; // Whether we should skip the incoming message since it could be a stale message.
   private boolean resultReusesInput;
@@ -29,13 +29,13 @@ public MergeConflictResult(
 
   public MergeConflictResult(
       ByteBuffer newValue,
-      Optional<GenericRecord> valueDeserialized,
+      Optional<GenericRecord> deserializedValue,
       int valueSchemaID,
       boolean resultReusesInput,
       GenericRecord rmdRecord) {
     this.updateIgnored = false;
     this.newValue = newValue;
-    this.valueDeserialized = valueDeserialized;
+    this.deserializedValue = deserializedValue;
     this.valueSchemaId = valueSchemaID;
     this.resultReusesInput = resultReusesInput;
     this.rmdRecord = rmdRecord;
@@ -75,7 +75,7 @@ public GenericRecord getRmdRecord() {
    * deserialize the value to generate the MCR.
    * @return deserialized new value if possible.
    */
-  public Optional<GenericRecord> getValueDeserialized() {
-    return valueDeserialized;
+  public Optional<GenericRecord> getDeserializedValue() {
+    return deserializedValue;
   }
 }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/view/ChangeCaptureViewWriter.java

@@ -66,7 +66,7 @@ public CompletableFuture<Void> processRecord(
       int newValueSchemaId,
       int oldValueSchemaId,
       GenericRecord replicationMetadataRecord,
-      Lazy<GenericRecord> newValueProvider) {
+      Lazy<GenericRecord> valueProvider) {
     // TODO: not sold about having currentValue in the interface but it VASTLY simplifies a lot of things with regards
     // to dealing with compression/chunking/etc. in the storage layer.
 

clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/view/MaterializedViewWriter.java

@@ -65,8 +65,8 @@ public CompletableFuture<Void> processRecord(
       int newValueSchemaId,
       int oldValueSchemaId,
       GenericRecord replicationMetadataRecord,
-      Lazy<GenericRecord> newValueProvider) {
-    return processRecord(newValue, key, newValueSchemaId, false, newValueProvider);
+      Lazy<GenericRecord> valueProvider) {
+    return processRecord(newValue, key, newValueSchemaId, false, valueProvider);
   }
 
   /**

clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/view/VeniceViewWriter.java

@@ -65,7 +65,8 @@ public VeniceViewWriter(
    * @param newValueSchemaId the schemaId of the incoming record
    * @param oldValueSchemaId the schemaId of the old record
    * @param replicationMetadataRecord the associated RMD for the incoming record.
-   * @param newValueProvider to provide the deserialized new value
+   * @param valueProvider to provide the corresponding deserialized newValue for PUT and UPDATE or the old value for the
+   *                      given key for DELETE.
    */
   public abstract CompletableFuture<Void> processRecord(
       ByteBuffer newValue,
@@ -74,7 +75,7 @@ public abstract CompletableFuture<Void> processRecord(
       int newValueSchemaId,
       int oldValueSchemaId,
       GenericRecord replicationMetadataRecord,
-      Lazy<GenericRecord> newValueProvider);
+      Lazy<GenericRecord> valueProvider);
 
   /**
    * To be called as a given ingestion task consumes each record. This is called prior to writing to a

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

@@ -769,21 +769,6 @@ public void run() {
           pushJobSetting.rewindTimeInSecondsOverride = DEFAULT_RE_PUSH_REWIND_IN_SECONDS_OVERRIDE;
           LOGGER.info("Overriding re-push rewind time in seconds to: {}", pushJobSetting.rewindTimeInSecondsOverride);
         }
-        if (pushJobSetting.repushTTLEnabled) {
-          // Build the full path for HDFSRmdSchemaSource:
-          // RMD schemas: <job_temp_dir>/rmd_schemas
-          // Value schemas: <job_temp_dir>/value_schemas
-          Path rmdSchemaDir = new Path(jobTmpDir, "rmd_schemas");
-          HadoopUtils.createDirectoryWithPermission(rmdSchemaDir, PERMISSION_700);
-          Path valueSchemaDir = new Path(jobTmpDir, "value_schemas");
-          HadoopUtils.createDirectoryWithPermission(valueSchemaDir, PERMISSION_700);
-          try (HDFSSchemaSource schemaSource =
-              new HDFSSchemaSource(valueSchemaDir, rmdSchemaDir, pushJobSetting.storeName)) {
-            schemaSource.saveSchemasOnDisk(controllerClient);
-            pushJobSetting.rmdSchemaDir = schemaSource.getRmdSchemaPath();
-            pushJobSetting.valueSchemaDir = schemaSource.getValueSchemaPath();
-          }
-        }
       }
       // Create new store version, topic and fetch Kafka url from backend
       createNewStoreVersion(
@@ -825,6 +810,9 @@ public void run() {
       } else {
         // Populate any view configs to job properties
         configureJobPropertiesWithMaterializedViewConfigs();
+        if (pushJobSetting.repushTTLEnabled || pushJobSetting.materializedViewConfigFlatMap != null) {
+          buildHDFSSchemaDir();
+        }
         if (pushJobSetting.sendControlMessagesDirectly) {
           getVeniceWriter(pushJobSetting).broadcastStartOfPush(
               SORTED,
@@ -930,6 +918,21 @@ public void run() {
     }
   }
 
+  private void buildHDFSSchemaDir() throws IOException {
+    // Build the full path for HDFSRmdSchemaSource:
+    // RMD schemas: <job_temp_dir>/rmd_schemas
+    // Value schemas: <job_temp_dir>/value_schemas
+    Path rmdSchemaDir = new Path(jobTmpDir, "rmd_schemas");
+    HadoopUtils.createDirectoryWithPermission(rmdSchemaDir, PERMISSION_700);
+    Path valueSchemaDir = new Path(jobTmpDir, "value_schemas");
+    HadoopUtils.createDirectoryWithPermission(valueSchemaDir, PERMISSION_700);
+    try (HDFSSchemaSource schemaSource = new HDFSSchemaSource(valueSchemaDir, rmdSchemaDir, pushJobSetting.storeName)) {
+      schemaSource.saveSchemasOnDisk(controllerClient);
+      pushJobSetting.rmdSchemaDir = schemaSource.getRmdSchemaPath();
+      pushJobSetting.valueSchemaDir = schemaSource.getValueSchemaPath();
+    }
+  }
+
   /**
    * Get the set of regions that haven't been pushed yet after targeted region push.
    * @return a set of regions that haven't been pushed yet.