[improve] improve batch iterator has next check (#164)

apache · Nov 27, 2023 · 3efe6ff · 3efe6ff
1 parent 86ca733
commit 3efe6ff
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/spark-doris-connector/src/main/scala/org/apache/doris/spark/writer/DorisWriter.scala b/spark-doris-connector/src/main/scala/org/apache/doris/spark/writer/DorisWriter.scala
@@ -150,7 +150,17 @@ class DorisWriter(settings: SparkSettings, preCommittedTxnAcc: CollectionAccumul
 
     private var isReset = false
 
-    override def hasNext: Boolean = recordCount < batchSize && iterator.hasNext
+    override def hasNext: Boolean = {
+      if (recordCount < batchSize) {
+        if (isReset) {
+          recordCount < buffer.size
+        } else {
+          iterator.hasNext
+        }
+      } else {
+        false
+      }
+    }
 
     override def next(): T = {
       recordCount += 1