This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit ffae3040b4b73f2e3dd82a9d87624e025ef68a74 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()
