This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new f778a9b1c9 [python] Skip empty batches when advancing merge readers
(#9820)
f778a9b1c9 is described below
commit f778a9b1c9c13e21ed99bbe5ca3207348d70b858
Author: zhigang <[email protected]>
AuthorDate: Tue Sep 15 11:18:50 2026 +0800
[python] Skip empty batches when advancing merge readers (#9820)
---
.../pypaimon/read/reader/sort_merge_reader.py | 23 +++---
.../pypaimon/tests/test_sort_merge_reader.py | 85 ++++++++++++++++++++++
2 files changed, 95 insertions(+), 13 deletions(-)
diff --git a/paimon-python/pypaimon/read/reader/sort_merge_reader.py
b/paimon-python/pypaimon/read/reader/sort_merge_reader.py
index c525a7592c..b1965b2795 100644
--- a/paimon-python/pypaimon/read/reader/sort_merge_reader.py
+++ b/paimon-python/pypaimon/read/reader/sort_merge_reader.py
@@ -149,19 +149,16 @@ class Element:
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:
diff --git a/paimon-python/pypaimon/tests/test_sort_merge_reader.py
b/paimon-python/pypaimon/tests/test_sort_merge_reader.py
new file mode 100644
index 0000000000..1dde5aa60e
--- /dev/null
+++ b/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()
