Skip to content

Forward iter_arrow, is_typed and features from RepeatExamplesIterable - #8395

Open
ebarkhordar wants to merge 1 commit into
huggingface:mainfrom
ebarkhordar:fix/repeat-iterable-arrow-properties
Open

Forward iter_arrow, is_typed and features from RepeatExamplesIterable#8395
ebarkhordar wants to merge 1 commit into
huggingface:mainfrom
ebarkhordar:fix/repeat-iterable-arrow-properties

Conversation

@ebarkhordar

Copy link
Copy Markdown
Contributor

RepeatExamplesIterable never forwarded iter_arrow, is_typed or features to the iterable it wraps, so .repeat() stranded an Arrow-capable child behind a non-Arrow parent: iteration fell back to the per-example path and the dtypes were re-inferred from Python objects.

This copies the three-property block and the _iter_arrow loop from SkipExamplesIterable and TakeExamplesIterable, which sit either side of it. _iter_arrow mirrors __iter__ exactly, including the per-pass reset of the child state dict, so state_dict()/load_state_dict() behaves the same on both paths.

Verification

On main at b7cb10b, Python 3.12, pyarrow 25.0.0, numpy 2.5.1:

  • The three new tests fail before the change and pass after. test_iterable_dataset_repeat_keeps_numpy_dtypes reports int64/float64/int64 before against a declared int32/float32/uint8; test_repeat_examples_iterable_arrow asserts iter_arrow is not None and runs both assert_load_state_dict_resumes_iteration and assert_load_state_dict_resumes_arrow_iteration.
  • tests/test_iterable_dataset.py goes from 441 passed to 445 passed, with the same 4 pre-existing test_interleave_dataset_with_sharding failures on both sides (a missing optional dependency in my environment, unrelated to this change).
  • Arrow-format iteration over 100k rows drops from 6.644s to 0.469s, which is parity with .take(100000) at 0.465s. list(ds.repeat(2)) is unchanged.
  • ruff format --check and ruff check are clean on both files.

What I did not check: the torch formatter, and non-scalar feature types such as Image. The dtype claim above is scoped to scalar Value columns under numpy format.

Fixes #8394

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

IterableDataset.repeat() drops the Arrow fast path and widens numpy dtypes

1 participant