feat(flink): Sort bulk insert records by record key for LSM layout

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsResolver.java

@@ -166,13 +166,13 @@ public static boolean isCowTable(Configuration conf) {
   /**
    * Returns the configured table storage layout.
    *
-   * <p>Insert and bulk insert operations preserve duplicate record keys and therefore default to
-   * the regular storage layout. Other operations use Flink's LSM tree default.
+   * <p>Insert operations preserve duplicate record keys and therefore default to the regular
+   * storage layout. Other operations use Flink's LSM tree default.
    */
   public static HoodieTableConfig.TableStorageLayout getTableStorageLayout(Configuration conf) {
     return HoodieTableConfig.TableStorageLayout.fromConfigValue(conf.getString(
         HoodieTableConfig.TABLE_STORAGE_LAYOUT.key(),
-        (isInsertOperation(conf) || isBulkInsertOperation(conf))
+        isInsertOperation(conf)
             ? HoodieTableConfig.TableStorageLayout.DEFAULT.configValue()
             : HoodieTableConfig.TableStorageLayout.LSM_TREE.configValue()));
   }

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java

@@ -47,7 +47,7 @@
 public class BucketBulkInsertWriterHelper extends BulkInsertWriterHelper {
   public static final String FILE_GROUP_META_FIELD = "_fg";
 
-  private final int recordArity;
+  protected final int recordArity;
 
   private String lastFileId; // for efficient code path
 
@@ -63,19 +63,27 @@ public void write(RowData tuple) throws IOException {
       String recordKey = keyGen.getRecordKey(record);
       String partitionPath = keyGen.getPartitionPath(record);
       String fileId = tuple.getString(0).toString();
-      if ((lastFileId == null) || !lastFileId.equals(fileId)) {
-        log.info("Creating new file for partition path {}", partitionPath);
-        handle = getRowCreateHandle(partitionPath, fileId);
-        lastFileId = fileId;
-      }
-      handle.write(recordKey, partitionPath, record);
+      writeRecord(recordKey, partitionPath, fileId, record);
     } catch (Throwable throwable) {
       IOException ioException = new IOException("Exception happened when bulk insert.", throwable);
       log.error("Global error thrown while trying to write records in HoodieRowDataCreateHandle", ioException);
       throw ioException;
     }
   }
 
+  protected void writeRecord(
+      String recordKey,
+      String partitionPath,
+      String fileId,
+      RowData record) throws IOException {
+    if ((lastFileId == null) || !lastFileId.equals(fileId)) {
+      log.info("Creating new file for partition path {}", partitionPath);
+      handle = getRowCreateHandle(partitionPath, fileId);
+      lastFileId = fileId;
+    }
+    handle.write(recordKey, partitionPath, record);
+  }
+
   private HoodieRowDataCreateHandle getRowCreateHandle(String partitionPath, String fileId) throws IOException {
     if (!handles.containsKey(fileId)) { // if there is no handle corresponding to the fileId
       if (this.isInputSorted) {
@@ -93,19 +101,30 @@ public static SortOperatorGen getFileIdSorterGen(RowType rowType) {
     return new SortOperatorGen(rowType, new String[] {FILE_GROUP_META_FIELD});
   }
 
-  private static String getFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, List<String> indexKeyFields,
-                                  NumBucketsFunction numBucketsFunction, boolean needFixedFileIdSuffix) {
-    String recordKey = keyGen.getRecordKey(record);
-    String partition = keyGen.getPartitionPath(record);
-    final int numBuckets = numBucketsFunction.getNumBuckets(partition);
+  static String getFileId(
+      Map<String, String> bucketIdToFileId,
+      String recordKey,
+      String partitionPath,
+      List<String> indexKeyFields,
+      NumBucketsFunction numBucketsFunction,
+      boolean needFixedFileIdSuffix) {
+    final int numBuckets = numBucketsFunction.getNumBuckets(partitionPath);
     final int bucketNum = BucketIdentifier.getBucketId(recordKey, indexKeyFields, numBuckets);
-    String bucketId = partition + bucketNum;
+    String bucketId = partitionPath + bucketNum;
     return bucketIdToFileId.computeIfAbsent(bucketId, k -> needFixedFileIdSuffix ? BucketIdentifier.newBucketFileIdForNBCC(bucketNum) : BucketIdentifier.newBucketFileIdPrefix(bucketNum));
   }
 
   public static RowData rowWithFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, List<String> indexKeyFields,
                                       NumBucketsFunction numBucketsFunction, boolean needFixedFileIdSuffix) {
-    final String fileId = getFileId(bucketIdToFileId, keyGen, record, indexKeyFields, numBucketsFunction, needFixedFileIdSuffix);
+    String recordKey = keyGen.getRecordKey(record);
+    String partitionPath = keyGen.getPartitionPath(record);
+    final String fileId = getFileId(
+        bucketIdToFileId,
+        recordKey,
+        partitionPath,
+        indexKeyFields,
+        numBucketsFunction,
+        needFixedFileIdSuffix);
     return GenericRowData.of(StringData.fromString(fileId), record);
   }
 

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/LsmBucketBulkInsertWriterHelper.java

@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.sink.bucket;
+
+import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.index.bucket.partition.NumBucketsFunction;
+import org.apache.hudi.sink.bulk.RowDataKeyGen;
+import org.apache.hudi.table.HoodieTable;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.table.types.logical.RowType;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Bucket-index bulk-insert writer helper for LSM input rows.
+ *
+ * <p>The input row contains file ID, encoded record key, and the original table row.
+ */
+public class LsmBucketBulkInsertWriterHelper extends BucketBulkInsertWriterHelper {
+
+  public LsmBucketBulkInsertWriterHelper(
+      Configuration conf,
+      HoodieTable<?, ?, ?, ?> hoodieTable,
+      HoodieWriteConfig writeConfig,
+      String instantTime,
+      int taskPartitionId,
+      long taskId,
+      long taskEpochId,
+      RowType rowType) {
+    super(conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType);
+  }
+
+  @Override
+  public void write(RowData sortRow) throws IOException {
+    String fileId = sortRow.getString(0).toString();
+    String recordKey = sortRow.getString(1).toString();
+    RowData record = sortRow.getRow(2, recordArity);
+    String partitionPath = keyGen.getPartitionPath(record);
+    writeRecord(recordKey, partitionPath, fileId, record);
+  }
+
+  public static RowData rowWithFileIdAndRecordKey(
+      Map<String, String> bucketIdToFileId,
+      RowDataKeyGen keyGen,
+      RowData record,
+      List<String> indexKeyFields,
+      NumBucketsFunction numBucketsFunction,
+      boolean needFixedFileIdSuffix) {
+    String recordKey = keyGen.getRecordKey(record);
+    String partitionPath = keyGen.getPartitionPath(record);
+    String fileId = getFileId(
+        bucketIdToFileId,
+        recordKey,
+        partitionPath,
+        indexKeyFields,
+        numBucketsFunction,
+        needFixedFileIdSuffix);
+    return GenericRowData.of(
+        StringData.fromString(fileId),
+        StringData.fromString(recordKey),
+        record);
+  }
+}

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java

@@ -115,7 +115,9 @@ public BulkInsertWriterHelper(Configuration conf, HoodieTable hoodieTable, Hoodi
         ? schema
         : HoodieSchemaUtils.addMetadataFields(schema, writeConfig.allowOperationMetadataField());
     this.preserveHoodieMetadata = preserveHoodieMetadata;
-    this.isInputSorted = OptionsResolver.isBulkInsertOperation(conf) && conf.get(FlinkOptions.WRITE_BULK_INSERT_SORT_INPUT);
+    this.isInputSorted = OptionsResolver.isBulkInsertOperation(conf)
+        && (conf.get(FlinkOptions.WRITE_BULK_INSERT_SORT_INPUT)
+        || OptionsResolver.isLsmTreeStorageLayout(conf));
     this.fileIdPrefix = UUID.randomUUID().toString();
     this.keyGen = preserveHoodieMetadata ? null : RowDataKeyGens.instance(conf, rowType, taskPartitionId, instantTime);
     this.writeMetrics = writeMetrics;
@@ -129,21 +131,26 @@ public void write(RowData record) throws IOException {
       String partitionPath = preserveHoodieMetadata
           ? record.getString(HoodieRecord.PARTITION_PATH_META_FIELD_ORD).toString()
           : keyGen.getPartitionPath(record);
-
-      if ((lastKnownPartitionPath == null) || !lastKnownPartitionPath.equals(partitionPath) || !handle.canWrite()) {
-        handle = getRowCreateHandle(partitionPath);
-        lastKnownPartitionPath = partitionPath;
-        writeMetrics.ifPresent(FlinkStreamWriteMetrics::markHandleSwitch);
-      }
-      handle.write(recordKey, partitionPath, record);
-      writeMetrics.ifPresent(FlinkStreamWriteMetrics::markRecordIn);
+      writeRecord(recordKey, partitionPath, record);
     } catch (Throwable t) {
       IOException ioException = new IOException("Exception happened when bulk insert.", t);
       log.error("Global error thrown while trying to write records in HoodieRowCreateHandle ", ioException);
       throw new IOException(ioException);
     }
   }
 
+  protected void writeRecord(String recordKey, String partitionPath, RowData record) throws IOException {
+    if ((lastKnownPartitionPath == null)
+        || !lastKnownPartitionPath.equals(partitionPath)
+        || !handle.canWrite()) {
+      handle = getRowCreateHandle(partitionPath);
+      lastKnownPartitionPath = partitionPath;
+      writeMetrics.ifPresent(FlinkStreamWriteMetrics::markHandleSwitch);
+    }
+    handle.write(recordKey, partitionPath, record);
+    writeMetrics.ifPresent(FlinkStreamWriteMetrics::markRecordIn);
+  }
+
   private HoodieRowDataCreateHandle 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
@@ -228,4 +235,3 @@ private WriteStatus closeWriteHandle(HoodieRowDataCreateHandle rowCreateHandle)
   }
 
 }
-

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/LsmBulkInsertWriterHelper.java

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.sink.bulk;
+
+import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.table.HoodieTable;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.table.types.logical.RowType;
+
+import java.io.IOException;
+
+/**
+ * Bulk-insert writer helper for LSM input rows.
+ *
+ * <p>The input row contains partition path, encoded record key, and the original table row. The
+ * routing and record keys are retained after sorting so the writer does not generate them again.
+ */
+public class LsmBulkInsertWriterHelper extends BulkInsertWriterHelper {
+
+  private final int recordArity;
+
+  public LsmBulkInsertWriterHelper(
+      Configuration conf,
+      HoodieTable<?, ?, ?, ?> hoodieTable,
+      HoodieWriteConfig writeConfig,
+      String instantTime,
+      int taskPartitionId,
+      long taskId,
+      long taskEpochId,
+      RowType rowType) {
+    super(conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType);
+    this.recordArity = rowType.getFieldCount();
+  }
+
+  @Override
+  public void write(RowData sortRow) throws IOException {
+    String partitionPath = sortRow.getString(0).toString();
+    String recordKey = sortRow.getString(1).toString();
+    RowData record = sortRow.getRow(2, recordArity);
+    writeRecord(recordKey, partitionPath, record);
+  }
+
+  /**
+   * Decorates a table row with the partition path and encoded Hudi record key used by the LSM
+   * sorter.
+   *
+   * @param partitionPath partition path
+   * @param record        original table row used to generate the Hudi record key
+   * @param keyGen        Hudi RowData key generator
+   * @return internal LSM sort row containing partition path, record key, and original table row
+   */
+  public static RowData rowWithPartitionAndKey(
+      String partitionPath,
+      RowData record,
+      RowDataKeyGen keyGen) {
+    return GenericRowData.of(
+        StringData.fromString(partitionPath),
+        StringData.fromString(keyGen.getRecordKey(record)),
+        record);
+  }
+}

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/WriterHelpers.java

@@ -21,6 +21,7 @@
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.configuration.OptionsResolver;
 import org.apache.hudi.sink.bucket.BucketBulkInsertWriterHelper;
+import org.apache.hudi.sink.bucket.LsmBucketBulkInsertWriterHelper;
 import org.apache.hudi.table.HoodieTable;
 
 import org.apache.flink.configuration.Configuration;
@@ -32,8 +33,18 @@
 public class WriterHelpers {
   public static BulkInsertWriterHelper getWriterHelper(Configuration conf, HoodieTable<?, ?, ?, ?> hoodieTable, HoodieWriteConfig writeConfig,
                                                        String instantTime, int taskPartitionId, long taskId, long taskEpochId, RowType rowType) {
-    return OptionsResolver.isBucketIndexType(conf)
-        ? new BucketBulkInsertWriterHelper(conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType)
-        : new BulkInsertWriterHelper(conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType);
+    if (OptionsResolver.isLsmTreeStorageLayout(conf)) {
+      return OptionsResolver.isBucketIndexType(conf)
+          ? new LsmBucketBulkInsertWriterHelper(
+              conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType)
+          : new LsmBulkInsertWriterHelper(
+              conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType);
+    } else {
+      return OptionsResolver.isBucketIndexType(conf)
+          ? new BucketBulkInsertWriterHelper(
+          conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType)
+          : new BulkInsertWriterHelper(
+              conf, hoodieTable, writeConfig, instantTime, taskPartitionId, taskId, taskEpochId, rowType);
+    }
   }
 }