[HUDI-9591] FG reader based merge handle for COW merge

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java

@@ -73,7 +73,6 @@
 import org.apache.hudi.index.HoodieIndex;
 import org.apache.hudi.io.FileGroupReaderBasedMergeHandle;
 import org.apache.hudi.io.HoodieConcatHandle;
-import org.apache.hudi.io.HoodieWriteMergeHandle;
 import org.apache.hudi.keygen.SimpleAvroKeyGenerator;
 import org.apache.hudi.keygen.constant.KeyGeneratorOptions;
 import org.apache.hudi.keygen.constant.KeyGeneratorType;
@@ -855,7 +854,7 @@ public class HoodieWriteConfig extends HoodieConfig {
 
   public static final ConfigProperty<String> MERGE_HANDLE_CLASS_NAME = ConfigProperty
       .key("hoodie.write.merge.handle.class")
-      .defaultValue(HoodieWriteMergeHandle.class.getName())
+      .defaultValue(FileGroupReaderBasedMergeHandle.class.getName())
       .markAdvanced()
       .sinceVersion("1.1.0")
       .withDocumentation("The merge handle class that implements interface{@link HoodieMergeHandle} to merge the records "

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/io/HoodieMergeHandleFactory.java

@@ -124,12 +124,10 @@ public static <T, I, K, O> HoodieMergeHandle<T, I, K, O> create(
     String mergeHandleClass = config.getCompactionMergeHandleClassName();
     String logContext = String.format("for fileId %s and partitionPath %s at commit %s", operation.getFileId(), operation.getPartitionPath(), instantTime);
     LOG.info("Create HoodieMergeHandle implementation {} {}", mergeHandleClass, logContext);
-
     Class<?>[] constructorParamTypes = new Class<?>[] {
         HoodieWriteConfig.class, String.class, HoodieTable.class, CompactionOperation.class,
         TaskContextSupplier.class, HoodieReaderContext.class, String.class, HoodieRecord.HoodieRecordType.class
     };
-
     return instantiateMergeHandle(
         isFallbackEnabled, mergeHandleClass, COMPACT_MERGE_HANDLE_CLASS_NAME.defaultValue(), logContext, constructorParamTypes,
         config, instantTime, hoodieTable, operation, taskContextSupplier, readerContext, maxInstantTime, recordType);

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/io/HoodieWriteMergeHandle.java

@@ -429,7 +429,7 @@ protected void writeIncomingRecords() throws IOException {
     }
   }
 
-  private Schema getNewSchema() {
+  protected Schema getNewSchema() {
     return preserveMetadata ? writeSchemaWithMetaFields : writeSchema;
   }
 

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/io/IOUtils.java

@@ -119,6 +119,7 @@ public static Iterator<List<WriteStatus>> runMerge(HoodieMergeHandle<?, ?, ?, ?>
           "Error in finding the old file path at commit " + instantTime + " for fileId: " + fileId);
     } else {
       mergeHandle.doMerge();
+      mergeHandle.close();
     }
 
     // TODO(vc): This needs to be revisited

hudi-client/hudi-java-client/src/main/java/org/apache/hudi/table/action/commit/BaseJavaCommitActionExecutor.java

@@ -21,6 +21,7 @@
 import org.apache.hudi.client.WriteStatus;
 import org.apache.hudi.common.data.HoodieListData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
+import org.apache.hudi.common.engine.HoodieReaderContext;
 import org.apache.hudi.common.model.HoodieKey;
 import org.apache.hudi.common.model.HoodieRecord;
 import org.apache.hudi.common.model.HoodieRecordLocation;
@@ -37,6 +38,7 @@
 import org.apache.hudi.exception.HoodieUpsertException;
 import org.apache.hudi.execution.JavaLazyInsertIterable;
 import org.apache.hudi.io.CreateHandleFactory;
+import org.apache.hudi.io.FileGroupReaderBasedMergeHandle;
 import org.apache.hudi.io.HoodieMergeHandle;
 import org.apache.hudi.io.HoodieMergeHandleFactory;
 import org.apache.hudi.io.IOUtils;
@@ -256,6 +258,12 @@ public Iterator<List<WriteStatus>> handleUpdate(String partitionPath, String fil
             + "columns are disabled. Please choose the right key generator if you wish to disable meta fields.", e);
       }
     }
+    if (config.getMergeHandleClassName().equals(FileGroupReaderBasedMergeHandle.class.getName())) {
+      HoodieReaderContext<T> readerContext = table.getContext().<T>getReaderContextFactory(table.getMetaClient()).getContext();
+      return new FileGroupReaderBasedMergeHandle<>(
+          config, instantTime, table, Option.of(recordItr), partitionPath, fileId,
+          taskContextSupplier, keyGeneratorOpt, readerContext, instantTime, config.getRecordMerger().getRecordType(), Option.empty());
+    }
     return HoodieMergeHandleFactory.create(operationType, config, instantTime, table, recordItr, partitionPath, fileId,
         taskContextSupplier, keyGeneratorOpt);
   }

hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BaseSparkCommitActionExecutor.java

@@ -19,6 +19,7 @@
 package org.apache.hudi.table.action.commit;
 
 import org.apache.hudi.client.utils.SparkPartitionUtils;
+import org.apache.hudi.common.engine.HoodieReaderContext;
 import org.apache.hudi.index.HoodieSparkIndexClient;
 import org.apache.hudi.client.WriteStatus;
 import org.apache.hudi.client.clustering.update.strategy.SparkAllowUpdateStrategy;
@@ -47,6 +48,7 @@
 import org.apache.hudi.exception.HoodieUpsertException;
 import org.apache.hudi.execution.SparkLazyInsertIterable;
 import org.apache.hudi.index.HoodieIndex;
+import org.apache.hudi.io.FileGroupReaderBasedMergeHandle;
 import org.apache.hudi.io.HoodieMergeHandle;
 import org.apache.hudi.io.CreateHandleFactory;
 import org.apache.hudi.io.HoodieMergeHandleFactory;
@@ -176,7 +178,7 @@ public HoodieWriteMetadata<HoodieData<WriteStatus>> execute(HoodieData<HoodieRec
       LOG.info("RDD PreppedRecords was persisted at: {}", inputRDD.getStorageLevel());
     }
 
-    // Handle records update with clustering
+    // Handle records update with clustering.
     HoodieData<HoodieRecord<T>> inputRecordsWithClusteringUpdate = clusteringHandleUpdate(inputRecords);
     LOG.info("Num spark partitions for inputRecords before triggering workload profile {}", inputRecordsWithClusteringUpdate.getNumPartitions());
 
@@ -387,8 +389,17 @@ public Iterator<List<WriteStatus>> handleUpdate(String partitionPath, String fil
   }
 
   protected HoodieMergeHandle getUpdateHandle(String partitionPath, String fileId, Iterator<HoodieRecord<T>> recordItr) {
-    HoodieMergeHandle mergeHandle = HoodieMergeHandleFactory.create(operationType, config, instantTime, table, recordItr, partitionPath, fileId,
-        taskContextSupplier, keyGeneratorOpt);
+    HoodieMergeHandle mergeHandle;
+    if (config.getMergeHandleClassName().equals(FileGroupReaderBasedMergeHandle.class.getName())) {
+      HoodieReaderContext<T> readerContext =
+          table.getContext().<T>getReaderContextFactory(table.getMetaClient()).getContext();
+      mergeHandle = new FileGroupReaderBasedMergeHandle(
+          config, instantTime, table, Option.of(recordItr), partitionPath, fileId, taskContextSupplier, keyGeneratorOpt,
+          readerContext, instantTime, table.getConfig().getRecordMerger().getRecordType(), Option.empty());
+    } else {
+      mergeHandle = HoodieMergeHandleFactory.create(
+          operationType, config, instantTime, table, recordItr, partitionPath, fileId, taskContextSupplier, keyGeneratorOpt);
+    }
     if (mergeHandle.getOldFilePath() != null && mergeHandle.baseFileForMerge().getBootstrapBaseFile().isPresent()) {
       Option<String[]> partitionFields = table.getMetaClient().getTableConfig().getPartitionFields();
       Object[] partitionValues = SparkPartitionUtils.getPartitionFieldVals(partitionFields, mergeHandle.getPartitionPath(),

hudi-common/src/main/java/org/apache/hudi/common/table/read/BaseFileUpdateCallback.java

@@ -44,4 +44,9 @@ public interface BaseFileUpdateCallback<T> {
    * @param previousRecord the record in the base file before deletion
    */
   void onDelete(String recordKey, T previousRecord);
+
+  /**
+   * Callback method to return the name of the callback.
+   */
+  String getName();
 }

hudi-common/src/main/java/org/apache/hudi/common/table/read/BufferedRecord.java

@@ -74,6 +74,20 @@ public static <T> BufferedRecord<T> forRecordWithContext(T record, Schema schema
     return new BufferedRecord<>(recordKey, orderingValue, record, schemaId, isDelete);
   }
 
+  /**
+   * When HoodieRecord is given, recordKey is definitely available.
+   */
+  public static <T> BufferedRecord<T> forRecordWithContext(String recordKey,
+                                                           T record,
+                                                           Schema schema,
+                                                           HoodieReaderContext<T> readerContext,
+                                                           Option<String> orderingFieldName,
+                                                           boolean isDelete) {
+    Integer schemaId = readerContext.encodeAvroSchema(schema);
+    Comparable orderingValue = readerContext.getOrderingValue(record, schema, orderingFieldName);
+    return new BufferedRecord<>(recordKey, orderingValue, record, schemaId, isDelete);
+  }
+
   public static <T> BufferedRecord<T> forDeleteRecord(DeleteRecord deleteRecord, Comparable orderingValue) {
     return new BufferedRecord<>(deleteRecord.getRecordKey(), orderingValue, null, null, true);
   }

hudi-common/src/main/java/org/apache/hudi/common/table/read/HoodieFileGroupReader.java

@@ -55,6 +55,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
@@ -74,7 +75,7 @@
 public final class HoodieFileGroupReader<T> implements Closeable {
   private final HoodieReaderContext<T> readerContext;
   private final HoodieTableMetaClient metaClient;
-  private final InputSplit inputSplit;
+  private final InputSplit<T> inputSplit;
   private final Option<String[]> partitionPathFields;
   private final Option<String> orderingFieldName;
   private final HoodieStorage storage;
@@ -91,7 +92,7 @@ public final class HoodieFileGroupReader<T> implements Closeable {
   // considers the log records which are inflight.
   private final boolean allowInflightInstants;
   // Callback to run custom logic on updates to the base files for the file group
-  private final Option<BaseFileUpdateCallback> fileGroupUpdateCallback;
+  private final List<BaseFileUpdateCallback<T>> fileGroupUpdateCallbacks;
   private final boolean enableOptimizedLogBlockScan;
   // The list of instant times read from the log blocks, this value is used by the log-compaction to allow optimized log-block scans
   private List<String> validBlockInstants = Collections.emptyList();
@@ -109,16 +110,16 @@ public HoodieFileGroupReader(HoodieReaderContext<T> readerContext, HoodieStorage
       long start, long length, boolean shouldUseRecordPosition) {
     this(readerContext, storage, tablePath, latestCommitTime, dataSchema, requestedSchema, internalSchemaOpt,
         hoodieTableMetaClient, props, shouldUseRecordPosition, false, false, false,
-        InputSplit.fromFileSlice(fileSlice, start, length), Option.empty(), false);
+        InputSplit.fromFileSlice(fileSlice, start, length), Collections.emptyList(), false);
   }
 
   private HoodieFileGroupReader(HoodieReaderContext<T> readerContext, HoodieStorage storage, String tablePath,
                                 String latestCommitTime, Schema dataSchema, Schema requestedSchema,
                                 Option<InternalSchema> internalSchemaOpt, HoodieTableMetaClient hoodieTableMetaClient, TypedProperties props,
                                 boolean shouldUseRecordPosition, boolean allowInflightInstants, boolean emitDelete, boolean sortOutput,
-                                InputSplit inputSplit, Option<BaseFileUpdateCallback> updateCallback, boolean enableOptimizedLogBlockScan) {
+                                InputSplit<T> inputSplit, List<BaseFileUpdateCallback<T>> updateCallback, boolean enableOptimizedLogBlockScan) {
     this.readerContext = readerContext;
-    this.fileGroupUpdateCallback = updateCallback;
+    this.fileGroupUpdateCallbacks = updateCallback;
     this.metaClient = hoodieTableMetaClient;
     this.storage = storage;
     this.enableOptimizedLogBlockScan = enableOptimizedLogBlockScan;
@@ -154,7 +155,7 @@ private HoodieFileGroupReader(HoodieReaderContext<T> readerContext, HoodieStorag
     this.readStats = new HoodieReadStats();
     this.recordBuffer = getRecordBuffer(readerContext, hoodieTableMetaClient,
         readerContext.getMergeMode(), tableConfig.getPartialUpdateMode(), props,
-        isSkipMerge, shouldUseRecordPosition, readStats, emitDelete, sortOutput);
+        isSkipMerge, shouldUseRecordPosition, readStats, emitDelete, sortOutput, inputSplit.recordIterator);
     this.allowInflightInstants = allowInflightInstants;
   }
 
@@ -170,11 +171,17 @@ private FileGroupRecordBuffer<T> getRecordBuffer(HoodieReaderContext<T> readerCo
                                                    boolean shouldUseRecordPosition,
                                                    HoodieReadStats readStats,
                                                    boolean emitDelete,
-                                                   boolean sortOutput) {
+                                                   boolean sortOutput,
+                                                   Option<Iterator<HoodieRecord<T>>> inputRecordOpt) {
+    UpdateProcessor<T> updateProcessor = UpdateProcessor.create(readStats, readerContext, emitDelete, fileGroupUpdateCallbacks);
     if (inputSplit.logFiles.isEmpty()) {
+      if (inputRecordOpt.isPresent()) {
+        return new InputBasedFileGroupRecordBuffer<>(
+            readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode,
+            props, orderingFieldName, inputRecordOpt.get(), updateProcessor);
+      }
       return null;
     }
-    UpdateProcessor<T> updateProcessor = UpdateProcessor.create(readStats, readerContext, emitDelete, fileGroupUpdateCallback);
     if (isSkipMerge) {
       return new UnmergedFileGroupRecordBuffer<>(
           readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode, props, readStats);
@@ -183,7 +190,8 @@ private FileGroupRecordBuffer<T> getRecordBuffer(HoodieReaderContext<T> readerCo
           readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode, props, orderingFieldName, updateProcessor);
     } else if (shouldUseRecordPosition && inputSplit.baseFileOption.isPresent()) {
       return new PositionBasedFileGroupRecordBuffer<>(
-          readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode, inputSplit.baseFileOption.get().getCommitTime(), props, orderingFieldName, updateProcessor);
+          readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode, inputSplit.baseFileOption.get().getCommitTime(),
+          props, orderingFieldName, updateProcessor);
     } else {
       return new KeyBasedFileGroupRecordBuffer<>(
           readerContext, hoodieTableMetaClient, recordMergeMode, partialUpdateMode, props, orderingFieldName, updateProcessor);
@@ -196,7 +204,12 @@ private FileGroupRecordBuffer<T> getRecordBuffer(HoodieReaderContext<T> readerCo
   private void initRecordIterators() throws IOException {
     ClosableIterator<T> iter = makeBaseFileIterator();
     if (inputSplit.logFiles.isEmpty()) {
-      this.baseFileIterator = new CloseableMappingIterator<>(iter, readerContext::seal);
+      if (inputSplit.recordIterator.isPresent()) {
+        this.baseFileIterator = iter;
+        recordBuffer.setBaseFileIterator(baseFileIterator);
+      } else {
+        this.baseFileIterator = new CloseableMappingIterator<>(iter, readerContext::seal);
+      }
     } else {
       this.baseFileIterator = iter;
       scanLogFiles();
@@ -377,7 +390,8 @@ public ClosableIterator<T> getClosableIterator() throws IOException {
    */
   public ClosableIterator<HoodieRecord<T>> getClosableHoodieRecordIterator() throws IOException {
     return new CloseableMappingIterator<>(getClosableIterator(), nextRecord -> {
-      BufferedRecord<T> bufferedRecord = BufferedRecord.forRecordWithContext(nextRecord, readerContext.getSchemaHandler().getRequestedSchema(), readerContext, orderingFieldName, false);
+      BufferedRecord<T> bufferedRecord = BufferedRecord.forRecordWithContext(
+          nextRecord, readerContext.getSchemaHandler().getRequestedSchema(), readerContext, orderingFieldName, false);
       return readerContext.constructHoodieRecord(bufferedRecord);
     });
   }
@@ -454,7 +468,8 @@ public static class Builder<T> {
     private boolean emitDelete;
     private boolean sortOutput = false;
     private boolean enableOptimizedLogBlockScan = false;
-    private Option<BaseFileUpdateCallback> fileGroupUpdateCallback = Option.empty();
+    private List<BaseFileUpdateCallback<T>> fileGroupUpdateCallbacks = Collections.emptyList();
+    private Option<Iterator<HoodieRecord<T>>> recordIterator = Option.empty();
 
     public Builder<T> withReaderContext(HoodieReaderContext<T> readerContext) {
       this.readerContext = readerContext;
@@ -539,8 +554,8 @@ public Builder<T> withEmitDelete(boolean emitDelete) {
       return this;
     }
 
-    public Builder<T> withFileGroupUpdateCallback(Option<BaseFileUpdateCallback> fileGroupUpdateCallback) {
-      this.fileGroupUpdateCallback = fileGroupUpdateCallback;
+    public Builder<T> withFileGroupUpdateCallback(List<BaseFileUpdateCallback<T>> fileGroupUpdateCallbacks) {
+      this.fileGroupUpdateCallbacks = fileGroupUpdateCallbacks;
       return this;
     }
 
@@ -560,6 +575,11 @@ public Builder<T> withSortOutput(boolean sortOutput) {
       return this;
     }
 
+    public Builder<T> withRecordIterator(Option<Iterator<HoodieRecord<T>>> iterator) {
+      this.recordIterator = iterator;
+      return this;
+    }
+
     public HoodieFileGroupReader<T> build() {
       ValidationUtils.checkArgument(readerContext != null, "Reader context is required");
       ValidationUtils.checkArgument(hoodieTableMetaClient != null, "Hoodie table meta client is required");
@@ -576,21 +596,33 @@ public HoodieFileGroupReader<T> build() {
       ValidationUtils.checkArgument(logFiles != null, "Log files stream is required");
       ValidationUtils.checkArgument(partitionPath != null, "Partition path is required");
 
-      InputSplit inputSplit = new InputSplit(baseFileOption, logFiles, partitionPath, start, length);
+      InputSplit<T> inputSplit = new InputSplit<>(baseFileOption, logFiles, partitionPath, start, length, recordIterator);
       return new HoodieFileGroupReader<>(
           readerContext, storage, tablePath, latestCommitTime, dataSchema, requestedSchema, internalSchemaOpt, hoodieTableMetaClient,
-          props, shouldUseRecordPosition, allowInflightInstants, emitDelete, sortOutput, inputSplit, fileGroupUpdateCallback, enableOptimizedLogBlockScan);
+          props, shouldUseRecordPosition, allowInflightInstants, emitDelete, sortOutput, inputSplit, fileGroupUpdateCallbacks, enableOptimizedLogBlockScan);
     }
   }
 
-  private static class InputSplit {
+  private static class InputSplit<T> {
     private final Option<HoodieBaseFile> baseFileOption;
     private final List<HoodieLogFile> logFiles;
     private final String partitionPath;
     // Byte offset to start reading from the base file
     private final long start;
     // Length of bytes to read from the base file
     private final long length;
+    // For input record case.
+    private Option<Iterator<HoodieRecord<T>>> recordIterator;
+
+    InputSplit(Option<HoodieBaseFile> baseFileOption,
+               Stream<HoodieLogFile> logFiles,
+               String partitionPath,
+               long start,
+               long length,
+               Option<Iterator<HoodieRecord<T>>> recordIterator) {
+      this(baseFileOption, logFiles, partitionPath, start, length);
+      this.recordIterator = recordIterator;
+    }
 
     InputSplit(Option<HoodieBaseFile> baseFileOption, Stream<HoodieLogFile> logFiles, String partitionPath, long start, long length) {
       this.baseFileOption = baseFileOption;
@@ -603,8 +635,7 @@ private static class InputSplit {
     }
 
     static InputSplit fromFileSlice(FileSlice fileSlice, long start, long length) {
-      return new InputSplit(fileSlice.getBaseFile(), fileSlice.getLogFiles(), fileSlice.getPartitionPath(),
-          start, length);
+      return new InputSplit(fileSlice.getBaseFile(), fileSlice.getLogFiles(), fileSlice.getPartitionPath(), start, length);
     }
   }
 }