[python] Skip empty batches when advancing merge readers

paimon-python/pypaimon/read/reader/sort_merge_reader.py

@@ -149,19 +149,16 @@ def __init__(self, kv: KeyValue, iterator: RecordIterator[KeyValue], reader: Rec
         self.reader = reader
 
     def update(self) -> bool:
-        next_kv = self.iterator.next()
-        if next_kv is not None:
-            self.kv = next_kv
-            return True
-
-        self.iterator = self.reader.read_batch()
-        if self.iterator is None:
-            self.reader.close()
-            return False
-
-        next_kv_from_new_batch = self.iterator.next()
-        self.kv = next_kv_from_new_batch
-        return True
+        while True:
+            next_kv = self.iterator.next()
+            if next_kv is not None:
+                self.kv = next_kv
+                return True
+
+            self.iterator = self.reader.read_batch()
+            if self.iterator is None:
+                self.reader.close()
+                return False
 
 
 class HeapEntry:

paimon-python/pypaimon/tests/test_sort_merge_reader.py

@@ -0,0 +1,85 @@
+################################################################################
+#  Licensed to the Apache Software Foundation (ASF) under one
+#  or more contributor license agreements.  See the NOTICE file
+#  distributed with this work for additional information
+#  regarding copyright ownership.  The ASF licenses this file
+#  to you under the Apache License, Version 2.0 (the
+#  "License"); you may not use this file except in compliance
+#  with the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import unittest
+
+from pypaimon.read.reader.iface.record_iterator import RecordIterator
+from pypaimon.read.reader.iface.record_reader import RecordReader
+from pypaimon.read.reader.sort_merge_reader import SortMergeReaderWithMinHeap
+from pypaimon.schema.data_types import AtomicType, DataField
+from pypaimon.schema.table_schema import TableSchema
+from pypaimon.table.row.key_value import KeyValue
+from pypaimon.table.row.row_kind import RowKind
+
+
+class _ListIterator(RecordIterator):
+    def __init__(self, records):
+        self.records = iter(records)
+
+    def next(self):
+        return next(self.records, None)
+
+
+class _BatchReader(RecordReader):
+    def __init__(self, batches):
+        self.batches = iter(batches)
+        self.closed = False
+
+    def read_batch(self):
+        batch = next(self.batches, None)
+        return None if batch is None else _ListIterator(batch)
+
+    def close(self):
+        self.closed = True
+
+
+def _kv(key, sequence, value):
+    return KeyValue(1, 1).replace((key, sequence, RowKind.INSERT.value, value))
+
+
+class SortMergeReaderTest(unittest.TestCase):
+    def test_empty_batches_between_records_and_before_eof(self):
+        first = _BatchReader([
+            [_kv(1, 1, 10)], [], [],
+            [_kv(3, 1, 30), _kv(5, 1, 50)], [], [],
+        ])
+        second = _BatchReader([
+            [_kv(2, 2, 20), _kv(3, 2, 300), _kv(4, 2, 40)],
+        ])
+        schema = TableSchema(
+            fields=[DataField(0, 'id', AtomicType('INT'))],
+            primary_keys=['id'],
+        )
+        reader = SortMergeReaderWithMinHeap([first, second], schema)
+        try:
+            actual = []
+            while True:
+                batch = reader.read_batch()
+                if batch is None:
+                    break
+                while True:
+                    kv = batch.next()
+                    if kv is None:
+                        break
+                    actual.append((kv.key.get_field(0), kv.value.get_field(0)))
+
+            self.assertEqual(actual, [(1, 10), (2, 20), (3, 300), (4, 40), (5, 50)])
+            self.assertTrue(first.closed)
+            self.assertTrue(second.closed)
+        finally:
+            reader.close()