From 5cf4c93791c460cb35bdd53a5e0b512b964781a5 Mon Sep 17 00:00:00 2001
From: YeJunHao <41894543+leaves12138@users.noreply.github.com>
Date: Wed, 11 Dec 2024 15:14:03 +0800
Subject: [PATCH] [core] Check file size after write bundle of records (#4685)

---
 .../java/org/apache/paimon/io/RollingFileWriter.java     | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriter.java
index 109b7574304e..29b9223b9a37 100644
--- a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriter.java
+++ b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriter.java
@@ -64,10 +64,9 @@ public long targetFileSize() {
         return targetFileSize;
     }
 
-    @VisibleForTesting
-    boolean rollingFile() throws IOException {
+    private boolean rollingFile(boolean forceCheck) throws IOException {
         return currentWriter.reachTargetSize(
-                recordCount % CHECK_ROLLING_RECORD_CNT == 0, targetFileSize);
+                forceCheck || recordCount % CHECK_ROLLING_RECORD_CNT == 0, targetFileSize);
     }
 
     @Override
@@ -81,7 +80,7 @@ public void write(T row) throws IOException {
             currentWriter.write(row);
             recordCount += 1;
 
-            if (rollingFile()) {
+            if (rollingFile(false)) {
                 closeCurrentWriter();
             }
         } catch (Throwable e) {
@@ -105,7 +104,7 @@ public void writeBundle(BundleRecords bundle) throws IOException {
             currentWriter.writeBundle(bundle);
             recordCount += bundle.rowCount();
 
-            if (rollingFile()) {
+            if (rollingFile(true)) {
                 closeCurrentWriter();
             }
         } catch (Throwable e) {