feat: implement iter_arrow for skip, take and step iterables #7972
+209
−1
Conversation
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
This commit optimizes streaming operations by implementing _iter_arrow for SkipExamplesIterable, TakeExamplesIterable, and StepExamplesIterable. Key Changes: - Enabled fast batch-level processing for .skip(n) and .take(n) on streaming datasets. - Optimized distributed sharding (StepExamplesIterable) to use Arrow's .take() for picking multiple records from a batch simultaneously. - Updated _init_state_dict and load_state_dict to ensure seamless checkpointing while using Arrow iteration. Performance Impact: Users will see significant speedups when skipping or taking examples in streaming mode, as the dataset no longer needs to fallback to row-by-row Python dictionary conversion for these operations. Testing: Added 6 new unit tests to ests/test_iterable_dataset.py covering functional correctness and state resumption for all three iterable types.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
This commit optimizes streaming operations by implementing _iter_arrow for SkipExamplesIterable, TakeExamplesIterable, and StepExamplesIterable.
Key Changes:
.skip(n)and.take(n)on streaming datasets, bypassing slow row-by-row iteration..take()to extract multiple records from a batch simultaneously.Performance Impact:
Users will experience significant performance gains when skipping or taking examples in streaming mode. By staying in the "Arrow path" and avoiding Python dictionary conversions, data loading overhead is drastically reduced, especially for large-scale training jobs.
Testing:
Integrated 6 new unit tests into tests/test_iterable_dataset.py to verify:
stepusing Arrow iteration.