[python] Fix RecursionError when reading a shard or slice of a large file - #10113
Open
jackylee-ch wants to merge 1 commit into
Open
jackylee-ch wants to merge 1 commit into
jackylee-ch wants to merge 1 commit into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
ShardBatchReader.read_arrow_batchskips each out-of-range batch by recursively calling itself, so recursion depth grows one per skipped batch. Awith_shard/with_sliceread whose range sits deep in a data file then overflows the stack withRecursionError: the pyarrow reader yields 1024-row batches by default, so a slice starting ~1M rows in skips >1000 batches (Python's recursion limit). A slice covering only the head of a large file hits it too — the batches afterend_posare drained the same recursive way.The fix iterates over skipped batches with a
whileloop, matchingConcatBatchReaderandApplyDeletionVectorReader. Behavior is otherwise unchanged.Tests
shard_batch_reader_test: 2000 single-row batches, slice(1999, 2000)— fails on master withRecursionError, passes here. Two more cases cover range filtering ([2,5)) and the straddle branches ([2,9)over 4-row batches).Written with Claude Code; verification is mine.