From e2a0f3d9532d610f809f8644f49eaac2cbfe02fa Mon Sep 17 00:00:00 2001
From: wuwenchi <wuwenchihdu@hotmail.com>
Date: Tue, 26 Dec 2023 11:00:03 +0800
Subject: [PATCH] support readbatch

---
 .../apache/doris/paimon/PaimonJniScanner.java    | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/fe/be-java-extensions/paimon-scanner/src/main/java/org/apache/doris/paimon/PaimonJniScanner.java b/fe/be-java-extensions/paimon-scanner/src/main/java/org/apache/doris/paimon/PaimonJniScanner.java
index 237e06ad9c51e4..849f6fb67bb504 100644
--- a/fe/be-java-extensions/paimon-scanner/src/main/java/org/apache/doris/paimon/PaimonJniScanner.java
+++ b/fe/be-java-extensions/paimon-scanner/src/main/java/org/apache/doris/paimon/PaimonJniScanner.java
@@ -57,6 +57,7 @@ public class PaimonJniScanner extends JniScanner {
     private long dbId;
     private long tblId;
     private long lastUpdateTime;
+    private RecordReader.RecordIterator<InternalRow> recordIterator = null;
 
     public PaimonJniScanner(int batchSize, Map<String, String> params) {
         LOG.debug("params:{}", params);
@@ -133,18 +134,25 @@ public void close() throws IOException {
     protected int getNext() throws IOException {
         int rows = 0;
         try {
-            RecordReader.RecordIterator<InternalRow> batch;
-            while ((batch = reader.readBatch()) != null) {
+            if (recordIterator == null) {
+                recordIterator = reader.readBatch();
+            }
+
+            while (recordIterator != null) {
                 InternalRow record;
-                while ((record = batch.next()) != null) {
+                while ((record = recordIterator.next()) != null) {
                     columnValue.setOffsetRow(record);
                     for (int i = 0; i < fields.length; i++) {
                         columnValue.setIdx(i, types[i]);
                         appendData(i, columnValue);
                     }
                     rows++;
+                    if (rows >= batchSize) {
+                        return rows;
+                    }
                 }
-                batch.releaseBatch();
+                recordIterator.releaseBatch();
+                recordIterator = reader.readBatch();
             }
         } catch (IOException e) {
             LOG.warn("failed to getNext columnValue ", e);