From 3efe6ffca0bf033aadcb558b49c2db35f9de6480 Mon Sep 17 00:00:00 2001 From: gnehil Date: Mon, 27 Nov 2023 11:12:22 +0800 Subject: [PATCH] [improve] improve batch iterator has next check (#164) --- .../org/apache/doris/spark/writer/DorisWriter.scala | 12 +++++++++++- 1 file changed, 11 insertions(+), 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 index 04856712..53187f4b 100644 --- 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