[server][common][vpj] Introduce VeniceComplexPartitioner to materialized view

The change will not work if record is actually large and chunked. Proper chunking
support is needed and will be addressed in a separate PR.

1. Introduced VeniceComplexPartitioner which extends VenicePartitioner and offer
a new API to partition by value and provide possible one-to-many partition mapping.

2. Added value provider of type Lazy<GenericRecord> to VeniceViewWriter's processRecord
API to access deserialized value if needed. e.g. when a VeniceComplexPartitioner is
involved.

3. MergeConflictResult will now provide deserialized value in a best effort manner.
This is useful when we already deserialized the value for a partial update operation
so that the deserialized value can be provided directly to the materialized view writer.

4. Refactored VeniceWriter to expose an API to write to desired partition with new
DIV. This is only used by the new method writeWithComplexPartitioner for now to handle
the partitioning and writes of the same value to mulitple partitions. However, this newly
exposed API should also come handy when we build proper chunking support to forward chunks
to predetermined view topic partitions.

5. writeWithComplexPartitioner in VeniceWriter will re-chunk when writing to each partition.
This should be optimized when we build proper chunking support.

Author: xunyin8

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

@@ -669,6 +669,14 @@ protected void processMessageAndMaybeProduceToKafka(
         ByteBuffer oldValueBB = mergeConflictResultWrapper.getOldValueByteBufferProvider().get();
         int oldValueSchemaId =
             oldValueBB == null ? -1 : mergeConflictResultWrapper.getOldValueProvider().get().writerSchemaId();
+        Lazy<GenericRecord> newValueProvider;
+        if (mergeConflictResult.getNewValueDeserialized().isPresent()) {
+          newValueProvider = Lazy.of(() -> mergeConflictResult.getNewValueDeserialized().get());
+        } else {
+          newValueProvider = getNewValueProvider(
+              mergeConflictResultWrapper.getUpdatedValueBytes(),
+              mergeConflictResult.getValueSchemaId());
+        }
         queueUpVersionTopicWritesWithViewWriters(
             partitionConsumptionState,
             (viewWriter) -> viewWriter.processRecord(
@@ -677,7 +685,8 @@ protected void processMessageAndMaybeProduceToKafka(
                 keyBytes,
                 mergeConflictResult.getValueSchemaId(),
                 oldValueSchemaId,
-                mergeConflictResult.getRmdRecord()),
+                mergeConflictResult.getRmdRecord(),
+                newValueProvider),
             produceToVersionTopic);
       } else {
         // This function may modify the original record in KME and it is unsafe to use the payload from KME directly

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

@@ -202,7 +202,7 @@ public class LeaderFollowerStoreIngestionTask extends StoreIngestionTask {
   protected final Map<String, VeniceViewWriter> viewWriters;
   protected final boolean hasChangeCaptureView;
 
-  protected final AvroStoreDeserializerCache storeDeserializerCache;
+  protected final AvroStoreDeserializerCache<GenericRecord> storeDeserializerCache;
 
   private final AtomicLong lastSendIngestionHeartbeatTimestamp = new AtomicLong(0);
 
@@ -3383,9 +3383,11 @@ protected void processMessageAndMaybeProduceToKafka(
       Put newPut = writeComputeResultWrapper.getNewPut();
       // keys will be serialized with chunk suffix during pass-through mode in L/F NR if chunking is enabled
       boolean isChunkedKey = isChunked() && !partitionConsumptionState.isEndOfPushReceived();
+      Lazy<GenericRecord> newValueProvider = getNewValueProvider(newPut.putValue, newPut.schemaId);
       queueUpVersionTopicWritesWithViewWriters(
           partitionConsumptionState,
-          (viewWriter) -> viewWriter.processRecord(newPut.putValue, keyBytes, newPut.schemaId, isChunkedKey),
+          (viewWriter) -> viewWriter
+              .processRecord(newPut.putValue, keyBytes, newPut.schemaId, isChunkedKey, newValueProvider),
           produceToVersionTopic);
     } else {
       produceToVersionTopic.run();
@@ -4043,4 +4045,10 @@ <T> T databaseLookupWithConcurrencyLimit(Supplier<T> supplier) {
       return supplier.get();
     }
   }
+  protected Lazy<GenericRecord> getNewValueProvider(ByteBuffer newValue, int schemaId) {
+    if (newValue == null) {
+      return Lazy.of(() -> null);
+    }
+    return Lazy.of(() -> storeDeserializerCache.getDeserializer(schemaId, schemaId).deserialize(newValue));
+  }
 }

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

@@ -32,6 +32,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -276,7 +277,12 @@ public MergeConflictResult update(
     final ByteBuffer updatedValueBytes = updatedValueAndRmd.getValue() == null
         ? null
         : serializeMergedValueRecord(oldValueSchemaID, updatedValueAndRmd.getValue());
-    return new MergeConflictResult(updatedValueBytes, oldValueSchemaID, false, updatedValueAndRmd.getRmd());
+    return new MergeConflictResult(
+        updatedValueBytes,
+        Optional.of(updatedValueAndRmd.getValue()),
+        oldValueSchemaID,
+        false,
+        updatedValueAndRmd.getRmd());
   }
 
   private MergeConflictResult mergePutWithValueLevelTimestamp(

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

@@ -1,6 +1,7 @@
 package com.linkedin.davinci.replication.merge;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
 import org.apache.avro.generic.GenericRecord;
 
 
@@ -12,6 +13,7 @@ public class MergeConflictResult {
   private static final MergeConflictResult IGNORED_RESULT = new MergeConflictResult();
 
   private ByteBuffer newValue;
+  private Optional<GenericRecord> newValueDeserialized;
   private int valueSchemaId;
   private final boolean updateIgnored; // Whether we should skip the incoming message since it could be a stale message.
   private boolean resultReusesInput;
@@ -22,8 +24,18 @@ public MergeConflictResult(
       int valueSchemaID,
       boolean resultReusesInput,
       GenericRecord rmdRecord) {
+    this(newValue, Optional.empty(), valueSchemaID, resultReusesInput, rmdRecord);
+  }
+
+  public MergeConflictResult(
+      ByteBuffer newValue,
+      Optional<GenericRecord> newValueDeserialized,
+      int valueSchemaID,
+      boolean resultReusesInput,
+      GenericRecord rmdRecord) {
     this.updateIgnored = false;
     this.newValue = newValue;
+    this.newValueDeserialized = newValueDeserialized;
     this.valueSchemaId = valueSchemaID;
     this.resultReusesInput = resultReusesInput;
     this.rmdRecord = rmdRecord;
@@ -56,4 +68,14 @@ public boolean doesResultReuseInput() {
   public GenericRecord getRmdRecord() {
     return rmdRecord;
   }
+
+  /**
+   * Provide the deserialized new value on a best-effort approach. Meaning that it's acceptable to return an empty
+   * Optional. e.g. MergeConflictResult of full PUTs will not contain deserialized new value since we don't need to
+   * deserialize the value to generate the MCR.
+   * @return deserialized new value if possible.
+   */
+  public Optional<GenericRecord> getNewValueDeserialized() {
+    return newValueDeserialized;
+  }
 }

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

@@ -17,6 +17,7 @@
 import com.linkedin.venice.pubsub.api.PubSubProduceResult;
 import com.linkedin.venice.schema.rmd.RmdUtils;
 import com.linkedin.venice.utils.VeniceProperties;
+import com.linkedin.venice.utils.lazy.Lazy;
 import com.linkedin.venice.views.ChangeCaptureView;
 import com.linkedin.venice.writer.VeniceWriter;
 import com.linkedin.venice.writer.VeniceWriterFactory;
@@ -65,7 +66,8 @@ public CompletableFuture<PubSubProduceResult> processRecord(
       byte[] key,
       int newValueSchemaId,
       int oldValueSchemaId,
-      GenericRecord replicationMetadataRecord) {
+      GenericRecord replicationMetadataRecord,
+      Lazy<GenericRecord> newValueProvider) {
     // 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.
 
@@ -87,7 +89,8 @@ public CompletableFuture<PubSubProduceResult> processRecord(
       ByteBuffer newValue,
       byte[] key,
       int newValueSchemaId,
-      boolean isChunkedKey) {
+      boolean isChunkedKey,
+      Lazy<GenericRecord> newValueProvider) {
     // No op
     return CompletableFuture.completedFuture(null);
   }

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

@@ -6,6 +6,7 @@
 import com.linkedin.venice.kafka.protocol.KafkaMessageEnvelope;
 import com.linkedin.venice.message.KafkaKey;
 import com.linkedin.venice.meta.Version;
+import com.linkedin.venice.partitioner.VeniceComplexPartitioner;
 import com.linkedin.venice.pubsub.PubSubProducerAdapterFactory;
 import com.linkedin.venice.pubsub.api.PubSubProduceResult;
 import com.linkedin.venice.serialization.KeyWithChunkingSuffixSerializer;
@@ -33,6 +34,7 @@ public class MaterializedViewWriter extends VeniceViewWriter {
   private final String materializedViewTopicName;
   private Lazy<VeniceWriter> veniceWriter;
   private final KeyWithChunkingSuffixSerializer keyWithChunkingSuffixSerializer = new KeyWithChunkingSuffixSerializer();
+  private final VeniceComplexPartitioner complexPartitioner;
 
   public MaterializedViewWriter(
       VeniceConfigLoader props,
@@ -48,6 +50,11 @@ public MaterializedViewWriter(
     this.veniceWriter = Lazy.of(
         () -> new VeniceWriterFactory(props.getCombinedProperties().toProperties(), pubSubProducerAdapterFactory, null)
             .createVeniceWriter(buildWriterOptions()));
+    if (internalView.getViewPartitioner() instanceof VeniceComplexPartitioner) {
+      complexPartitioner = (VeniceComplexPartitioner) internalView.getViewPartitioner();
+    } else {
+      complexPartitioner = null;
+    }
   }
 
   /**
@@ -64,8 +71,9 @@ public CompletableFuture<PubSubProduceResult> processRecord(
       byte[] key,
       int newValueSchemaId,
       int oldValueSchemaId,
-      GenericRecord replicationMetadataRecord) {
-    return processRecord(newValue, key, newValueSchemaId, false);
+      GenericRecord replicationMetadataRecord,
+      Lazy<GenericRecord> newValueProvider) {
+    return processRecord(newValue, key, newValueSchemaId, false, newValueProvider);
   }
 
   /**
@@ -81,16 +89,29 @@ public CompletableFuture<PubSubProduceResult> processRecord(
       ByteBuffer newValue,
       byte[] key,
       int newValueSchemaId,
-      boolean isChunkedKey) {
+      boolean isChunkedKey,
+      Lazy<GenericRecord> newValueProvider) {
     byte[] viewTopicKey = key;
     if (isChunkedKey) {
       viewTopicKey = keyWithChunkingSuffixSerializer.getKeyFromChunkedKey(key);
     }
-    if (newValue == null) {
-      // this is a delete operation
-      return veniceWriter.get().delete(viewTopicKey, null);
+    byte[] newValueBytes = newValue == null ? null : ByteUtils.extractByteArray(newValue);
+    if (complexPartitioner != null) {
+      return veniceWriter.get()
+          .writeWithComplexPartitioner(
+              viewTopicKey,
+              newValueBytes,
+              newValueSchemaId,
+              newValueProvider,
+              complexPartitioner,
+              internalView.getViewPartitionCount());
+    } else {
+      if (newValue == null) {
+        // this is a delete operation
+        return veniceWriter.get().delete(viewTopicKey, null);
+      }
+      return veniceWriter.get().put(viewTopicKey, newValueBytes, newValueSchemaId);
     }
-    return veniceWriter.get().put(viewTopicKey, ByteUtils.extractByteArray(newValue), newValueSchemaId);
   }
 
   @Override

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

@@ -10,6 +10,7 @@
 import com.linkedin.venice.message.KafkaKey;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.pubsub.api.PubSubProduceResult;
+import com.linkedin.venice.utils.lazy.Lazy;
 import com.linkedin.venice.views.VeniceView;
 import com.linkedin.venice.writer.VeniceWriterOptions;
 import java.nio.ByteBuffer;
@@ -62,18 +63,19 @@ public VeniceViewWriter(
    * @param newValue the incoming fully specified value which hasn't yet been committed to Venice
    * @param oldValue the previous value which has already been locally committed to Venice for the given key
    * @param key the key of the record that designates newValue and oldValue
-   * @param version the version of the store taking this record
    * @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
    */
   public abstract CompletableFuture<PubSubProduceResult> processRecord(
       ByteBuffer newValue,
       ByteBuffer oldValue,
       byte[] key,
       int newValueSchemaId,
       int oldValueSchemaId,
-      GenericRecord replicationMetadataRecord);
+      GenericRecord replicationMetadataRecord,
+      Lazy<GenericRecord> newValueProvider);
 
   /**
    * To be called as a given ingestion task consumes each record. This is called prior to writing to a
@@ -83,12 +85,14 @@ public abstract CompletableFuture<PubSubProduceResult> processRecord(
    * @param key the key of the record that designates newValue and oldValue
    * @param newValueSchemaId the schemaId of the incoming record
    * @param isChunkedKey is the key already serialized with {@link com.linkedin.venice.serialization.KeyWithChunkingSuffixSerializer}
+   * @param newValueProvider to provide the deserialized new value
    */
   public abstract CompletableFuture<PubSubProduceResult> processRecord(
       ByteBuffer newValue,
       byte[] key,
       int newValueSchemaId,
-      boolean isChunkedKey);
+      boolean isChunkedKey,
+      Lazy<GenericRecord> newValueProvider);
 
   /**
    * Called when the server encounters a control message. There isn't (today) a strict ordering

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTaskTest.java

@@ -280,7 +280,8 @@ public void testQueueUpVersionTopicWritesWithViewWriters() throws InterruptedExc
     viewWriterMap.put("testView", materializedViewWriter);
     when(mockVeniceViewWriterFactory.buildStoreViewWriters(any(), anyInt(), any())).thenReturn(viewWriterMap);
     CompletableFuture<PubSubProduceResult> viewWriterFuture = new CompletableFuture<>();
-    when(materializedViewWriter.processRecord(any(), any(), anyInt(), anyBoolean())).thenReturn(viewWriterFuture);
+    when(materializedViewWriter.processRecord(any(), any(), anyInt(), anyBoolean(), any()))
+        .thenReturn(viewWriterFuture);
     setUp();
     WriteComputeResultWrapper mockResult = mock(WriteComputeResultWrapper.class);
     Put put = new Put();
@@ -291,12 +292,12 @@ public void testQueueUpVersionTopicWritesWithViewWriters() throws InterruptedExc
         .thenReturn(CompletableFuture.completedFuture(null));
     leaderFollowerStoreIngestionTask.queueUpVersionTopicWritesWithViewWriters(
         mockPartitionConsumptionState,
-        (viewWriter) -> viewWriter.processRecord(mock(ByteBuffer.class), new byte[1], 1, false),
+        (viewWriter) -> viewWriter.processRecord(mock(ByteBuffer.class), new byte[1], 1, false, Lazy.of(() -> null)),
         () -> writeToVersionTopic.set(true));
     verify(mockPartitionConsumptionState, times(1)).getLastVTProduceCallFuture();
     ArgumentCaptor<CompletableFuture> vtWriteFutureCaptor = ArgumentCaptor.forClass(CompletableFuture.class);
     verify(mockPartitionConsumptionState, times(1)).setLastVTProduceCallFuture(vtWriteFutureCaptor.capture());
-    verify(materializedViewWriter, times(1)).processRecord(any(), any(), anyInt(), anyBoolean());
+    verify(materializedViewWriter, times(1)).processRecord(any(), any(), anyInt(), anyBoolean(), any());
     verify(hostLevelIngestionStats, times(1)).recordViewProducerLatency(anyDouble());
     verify(hostLevelIngestionStats, never()).recordViewProducerAckLatency(anyDouble());
     assertFalse(writeToVersionTopic.get());

clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/view/ChangeCaptureViewWriterTest.java

@@ -24,6 +24,7 @@
 import com.linkedin.venice.pubsub.api.PubSubProduceResult;
 import com.linkedin.venice.schema.rmd.RmdSchemaGenerator;
 import com.linkedin.venice.utils.VeniceProperties;
+import com.linkedin.venice.utils.lazy.Lazy;
 import com.linkedin.venice.views.ChangeCaptureView;
 import com.linkedin.venice.writer.VeniceWriter;
 import com.linkedin.venice.writer.VeniceWriterOptions;
@@ -245,15 +246,21 @@ public void testProcessRecord() throws ExecutionException, InterruptedException
     rmdRecordWithValueLevelTimeStamp.put(TIMESTAMP_FIELD_NAME, 20L);
     rmdRecordWithValueLevelTimeStamp.put(REPLICATION_CHECKPOINT_VECTOR_FIELD_NAME, vectors);
     changeCaptureViewWriter.setVeniceWriter(mockVeniceWriter);
-
+    Lazy<GenericRecord> dummyValueProvider = Lazy.of(() -> null);
     // Update Case
-    changeCaptureViewWriter.processRecord(NEW_VALUE, OLD_VALUE, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp).get();
+    changeCaptureViewWriter
+        .processRecord(NEW_VALUE, OLD_VALUE, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp, dummyValueProvider)
+        .get();
 
     // Insert Case
-    changeCaptureViewWriter.processRecord(NEW_VALUE, null, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp).get();
+    changeCaptureViewWriter
+        .processRecord(NEW_VALUE, null, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp, dummyValueProvider)
+        .get();
 
     // Deletion Case
-    changeCaptureViewWriter.processRecord(null, OLD_VALUE, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp).get();
+    changeCaptureViewWriter
+        .processRecord(null, OLD_VALUE, KEY, 1, 1, rmdRecordWithValueLevelTimeStamp, dummyValueProvider)
+        .get();
 
     // Set up argument captors
     ArgumentCaptor<byte[]> keyCaptor = ArgumentCaptor.forClass(byte[].class);

clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/view/MaterializedViewWriterTest.java

@@ -33,6 +33,7 @@
 import com.linkedin.venice.serialization.KeyWithChunkingSuffixSerializer;
 import com.linkedin.venice.utils.ObjectMapperFactory;
 import com.linkedin.venice.utils.VeniceProperties;
+import com.linkedin.venice.utils.lazy.Lazy;
 import com.linkedin.venice.views.MaterializedView;
 import com.linkedin.venice.views.VeniceView;
 import com.linkedin.venice.writer.VeniceWriter;
@@ -163,8 +164,12 @@ public void testViewWriterCanForwardChunkedKeysCorrectly() {
     for (int i = 0; i < 100; i++) {
       byte[] key = new byte[keySize];
       RANDOM.nextBytes(key);
-      materializedViewWriter
-          .processRecord(dummyValue, keyWithChunkingSuffixSerializer.serializeNonChunkedKey(key), 1, true);
+      materializedViewWriter.processRecord(
+          dummyValue,
+          keyWithChunkingSuffixSerializer.serializeNonChunkedKey(key),
+          1,
+          true,
+          Lazy.of(() -> null));
       verify(veniceWriter, times(1)).put(eq(key), any(), eq(1));
       Mockito.clearInvocations(veniceWriter);
     }

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/mapreduce/datawriter/jobs/DataWriterMRJob.java

@@ -46,6 +46,7 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.VALUE_FIELD_PROP;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.VALUE_SCHEMA_DIR;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.VALUE_SCHEMA_ID_PROP;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.VALUE_SCHEMA_STRING_PROP;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.VSON_PUSH;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.ZSTD_COMPRESSION_LEVEL;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.ZSTD_DICTIONARY_CREATION_SUCCESS;
@@ -272,6 +273,7 @@ protected void setupInputFormatConf(JobConf jobConf, PushJobSetting pushJobSetti
       if (pushJobSetting.isAvro) {
         jobConf.set(SCHEMA_STRING_PROP, pushJobSetting.inputDataSchemaString);
         jobConf.set(AvroJob.INPUT_SCHEMA, pushJobSetting.inputDataSchemaString);
+        jobConf.set(VALUE_SCHEMA_STRING_PROP, pushJobSetting.valueSchemaString);
         if (pushJobSetting.generatePartialUpdateRecordFromInput) {
           jobConf.setBoolean(GENERATE_PARTIAL_UPDATE_RECORD_FROM_INPUT, true);
           jobConf.set(UPDATE_SCHEMA_STRING_PROP, pushJobSetting.valueSchemaString);