From 650c4455c600b0346fed8b5b6aa4cc0bf3452e8c Mon Sep 17 00:00:00 2001
From: wangxianghu <wangxianghu@apache.org>
Date: Tue, 6 Jul 2021 00:30:57 +0800
Subject: [PATCH] [HUDI-2122] Improvement in packaging insert into smallfiles
 (#3213)

---
 .../apache/hudi/table/action/commit/UpsertPartitioner.java  | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java
index 9d60cde69..3ac81510b 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java
@@ -187,7 +187,7 @@ public class UpsertPartitioner<T extends HoodieRecordPayload<T>> extends Partiti
         for (SmallFile smallFile : smallFiles) {
           long recordsToAppend = Math.min((config.getParquetMaxFileSize() - smallFile.sizeBytes) / averageRecordSize,
               totalUnassignedInserts);
-          if (recordsToAppend > 0 && totalUnassignedInserts > 0) {
+          if (recordsToAppend > 0) {
             // create a new bucket or re-use an existing bucket
             int bucket;
             if (updateLocationToBucket.containsKey(smallFile.location.getFileId())) {
@@ -200,6 +200,10 @@ public class UpsertPartitioner<T extends HoodieRecordPayload<T>> extends Partiti
             bucketNumbers.add(bucket);
             recordsPerBucket.add(recordsToAppend);
             totalUnassignedInserts -= recordsToAppend;
+            if (totalUnassignedInserts <= 0) {
+              // stop the loop when all the inserts are assigned
+              break;
+            }
           }
         }