From 8860c549d5d5fff419187929f7ab0ab4d717e69b Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Tue, 22 Sep 2026 21:59:07 +0800 Subject: [PATCH] [python] Fix RecursionError when reading a shard or slice of a large file ShardBatchReader.read_arrow_batch skipped each out-of-range batch by recursively calling itself, so recursion depth grows one per skipped batch. A with_shard/with_slice read whose range sits deep in a data file then overflows the stack with RecursionError: the pyarrow reader yields 1024-row batches by default, so a slice starting ~1M rows into a file skips more than 1000 batches (Python's default recursion limit). A slice covering only the head of a large file hits it too, because the trailing batches after end_pos are drained the same recursive way before the reader returns None. Iterate over skipped batches with a while loop instead, matching the pattern already used by ConcatBatchReader and ApplyDeletionVectorReader. Behavior is otherwise unchanged. --- .../read/reader/shard_batch_reader.py | 13 ++- .../pypaimon/tests/shard_batch_reader_test.py | 94 +++++++++++++++++++ 2 files changed, 103 insertions(+), 4 deletions(-) create mode 100644 paimon-python/pypaimon/tests/shard_batch_reader_test.py diff --git a/paimon-python/pypaimon/read/reader/shard_batch_reader.py b/paimon-python/pypaimon/read/reader/shard_batch_reader.py index aa00814b82ea..445ff5f802d9 100644 --- a/paimon-python/pypaimon/read/reader/shard_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/shard_batch_reader.py @@ -38,8 +38,14 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: if isinstance(self.reader.format_reader, FormatBlobReader): # For blob reader, pass begin_idx and end_idx parameters return self.reader.read_arrow_batch(start_idx=self.start_pos, end_idx=self.end_pos) - else: - # For non-blob reader (DataFileBatchReader), use standard read_arrow_batch + + # For non-blob reader (DataFileBatchReader), use standard read_arrow_batch. + # Loop rather than recurse over skipped batches: a slice/shard whose range + # sits deep in a file (default parquet batch_size is 1024 rows) skips one + # batch per step, so recursing here overflows the stack (RecursionError) + # once the skipped count exceeds the interpreter limit. Mirrors the + # while-loop skip pattern in ConcatBatchReader / ApplyDeletionVectorReader. + while True: batch = self.reader.read_arrow_batch() if batch is None: @@ -56,8 +62,7 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: return batch.slice(self.start_pos - batch_begin, self.end_pos - self.start_pos) elif batch_begin < self.end_pos < self.current_pos: # batch ends after the desired range return batch.slice(0, self.end_pos - batch_begin) - else: # batch is outside the desired range - return self.read_arrow_batch() + # else: batch is outside the desired range -> read the next one (loop) def close(self): self.reader.close() diff --git a/paimon-python/pypaimon/tests/shard_batch_reader_test.py b/paimon-python/pypaimon/tests/shard_batch_reader_test.py new file mode 100644 index 000000000000..7dc4badbc298 --- /dev/null +++ b/paimon-python/pypaimon/tests/shard_batch_reader_test.py @@ -0,0 +1,94 @@ +# 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 + +import pyarrow as pa + +from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader +from pypaimon.read.reader.shard_batch_reader import ShardBatchReader + + +class _BatchReader(RecordBatchReader): + """A non-blob reader that replays an explicit list of arrow batches.""" + + format_reader = None # not a FormatBlobReader -> ShardBatchReader takes the row-range path + + def __init__(self, batches): + self._batches = iter(batches) + + def read_arrow_batch(self): + return next(self._batches, None) + + def close(self): + pass + + +def _single_row_batches(count): + return [pa.record_batch([pa.array([i])], names=["id"]) for i in range(count)] + + +def _read_all(reader): + got = [] + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + got.extend(batch.column("id").to_pylist()) + return got + + +class ShardBatchReaderTest(unittest.TestCase): + + def test_slice_deep_into_file_does_not_recurse(self): + # A slice/shard whose range sits many batches into a file must not recurse + # once per skipped batch. With the default parquet batch_size of 1024 rows a + # slice starting ~1M rows in skips >1000 batches; recursing there overflowed + # the stack with RecursionError. 2000 single-row batches reproduce that. + batch_count = 2000 + reader = ShardBatchReader( + _BatchReader(_single_row_batches(batch_count)), batch_count - 1, batch_count) + + batch = reader.read_arrow_batch() + + self.assertIsNotNone(batch) + self.assertEqual(batch.column("id").to_pylist(), [batch_count - 1]) + self.assertIsNone(reader.read_arrow_batch()) + + def test_slice_returns_only_rows_in_range(self): + # Semantics guard: with single-row batches, slice [2, 5) yields rows 2, 3, 4 + # and nothing else, so the loop refactor preserves the range filtering. + reader = ShardBatchReader(_BatchReader(_single_row_batches(8)), 2, 5) + + self.assertEqual(_read_all(reader), [2, 3, 4]) + + def test_slice_straddling_batch_boundaries(self): + # Multi-row batches so the two slice() branches are exercised: the first + # batch straddles start_pos (2 in [0,4)) and the last straddles end_pos + # (9 in [8,12)); slice [2, 9) must yield exactly rows 2..8. + batches = [ + pa.record_batch([pa.array([0, 1, 2, 3])], names=["id"]), + pa.record_batch([pa.array([4, 5, 6, 7])], names=["id"]), + pa.record_batch([pa.array([8, 9, 10, 11])], names=["id"]), + ] + reader = ShardBatchReader(_BatchReader(batches), 2, 9) + + self.assertEqual(_read_all(reader), [2, 3, 4, 5, 6, 7, 8]) + + +if __name__ == "__main__": + unittest.main()