[HUDI-1104] Adding support for UserDefinedPartitioners and SortModes to BulkInsert with Rows (#3149)

hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/DataSourceUtils.java

@@ -49,6 +49,8 @@ import org.apache.log4j.LogManager;
 import org.apache.log4j.Logger;
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -98,6 +100,25 @@ public class DataSourceUtils {
     }
   }
 
+  /**
+   * Create a UserDefinedBulkInsertPartitionerRows class via reflection,
+   * <br>
+   * if the class name of UserDefinedBulkInsertPartitioner is configured through the HoodieWriteConfig.
+   *
+   * @see HoodieWriteConfig#getUserDefinedBulkInsertPartitionerClass()
+   */
+  public static Option<BulkInsertPartitioner<Dataset<Row>>> createUserDefinedBulkInsertPartitionerWithRows(HoodieWriteConfig config)
+      throws HoodieException {
+    String bulkInsertPartitionerClass = config.getUserDefinedBulkInsertPartitionerClass();
+    try {
+      return StringUtils.isNullOrEmpty(bulkInsertPartitionerClass)
+          ? Option.empty() :
+          Option.of((BulkInsertPartitioner) ReflectionUtils.loadClass(bulkInsertPartitionerClass));
+    } catch (Throwable e) {
+      throw new HoodieException("Could not create UserDefinedBulkInsertPartitionerRows class " + bulkInsertPartitionerClass, e);
+    }
+  }
+
   /**
    * Create a payload class via reflection, passing in an ordering/precombine value.
    */

hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/internal/BulkInsertDataInternalWriterHelper.java

@@ -31,7 +31,9 @@ import org.apache.spark.sql.types.StructType;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 /**
@@ -48,15 +50,17 @@ public class BulkInsertDataInternalWriterHelper {
   private final HoodieTable hoodieTable;
   private final HoodieWriteConfig writeConfig;
   private final StructType structType;
+  private final Boolean arePartitionRecordsSorted;
   private final List<HoodieInternalWriteStatus> writeStatusList = new ArrayList<>();
 
   private HoodieRowCreateHandle handle;
   private String lastKnownPartitionPath = null;
   private String fileIdPrefix;
   private int numFilesWritten = 0;
+  private Map<String, HoodieRowCreateHandle> handles = new HashMap<>();
 
   public BulkInsertDataInternalWriterHelper(HoodieTable hoodieTable, HoodieWriteConfig writeConfig,
-      String instantTime, int taskPartitionId, long taskId, long taskEpochId, StructType structType) {
+      String instantTime, int taskPartitionId, long taskId, long taskEpochId, StructType structType, boolean arePartitionRecordsSorted) {
     this.hoodieTable = hoodieTable;
     this.writeConfig = writeConfig;
     this.instantTime = instantTime;
@@ -64,6 +68,7 @@ public class BulkInsertDataInternalWriterHelper {
     this.taskId = taskId;
     this.taskEpochId = taskEpochId;
     this.structType = structType;
+    this.arePartitionRecordsSorted = arePartitionRecordsSorted;
     this.fileIdPrefix = UUID.randomUUID().toString();
   }
 
@@ -74,7 +79,7 @@ public class BulkInsertDataInternalWriterHelper {
 
       if ((lastKnownPartitionPath == null) || !lastKnownPartitionPath.equals(partitionPath) || !handle.canWrite()) {
         LOG.info("Creating new file for partition path " + partitionPath);
-        createNewHandle(partitionPath);
+        handle = getRowCreateHandle(partitionPath);
         lastKnownPartitionPath = partitionPath;
       }
       handle.write(record);
@@ -92,19 +97,30 @@ public class BulkInsertDataInternalWriterHelper {
   public void abort() {
   }
 
-  private void createNewHandle(String partitionPath) throws IOException {
-    if (null != handle) {
-      close();
+  private HoodieRowCreateHandle getRowCreateHandle(String partitionPath) throws IOException {
+    if (!handles.containsKey(partitionPath)) { // if there is no handle corresponding to the partition path
+      // if records are sorted, we can close all existing handles
+      if (arePartitionRecordsSorted) {
+        close();
+      }
+      handles.put(partitionPath, new HoodieRowCreateHandle(hoodieTable, writeConfig, partitionPath, getNextFileId(),
+          instantTime, taskPartitionId, taskId, taskEpochId, structType));
+    } else if (!handles.get(partitionPath).canWrite()) {
+      // even if there is a handle to the partition path, it could have reached its max size threshold. So, we close the handle here and
+      // create a new one.
+      writeStatusList.add(handles.remove(partitionPath).close());
+      handles.put(partitionPath, new HoodieRowCreateHandle(hoodieTable, writeConfig, partitionPath, getNextFileId(),
+          instantTime, taskPartitionId, taskId, taskEpochId, structType));
     }
-    handle = new HoodieRowCreateHandle(hoodieTable, writeConfig, partitionPath, getNextFileId(),
-        instantTime, taskPartitionId, taskId, taskEpochId, structType);
+    return handles.get(partitionPath);
   }
 
   public void close() throws IOException {
-    if (null != handle) {
-      writeStatusList.add(handle.close());
-      handle = null;
+    for (HoodieRowCreateHandle rowCreateHandle: handles.values()) {
+      writeStatusList.add(rowCreateHandle.close());
     }
+    handles.clear();
+    handle = null;
   }
 
   private String getNextFileId() {