From abbd7f98cf0090afedd2e7eb356f85d5e4d35d9a Mon Sep 17 00:00:00 2001 From: connork Date: Tue, 13 Oct 2026 13:14:38 -0700 Subject: [PATCH] Skip complete Parquet sources using metadata counts --- src/datasets/iterable_dataset.py | 85 ++++- .../packaged_modules/parquet/parquet.py | 45 +++ tests/packaged_modules/test_parquet.py | 358 +++++++++++++++++- tests/test_iterable_dataset.py | 43 +++ 4 files changed, 525 insertions(+), 6 deletions(-) diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 23153de7920..4930817725e 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -391,10 +391,14 @@ def __init__( kwargs: dict, generate_more_kwargs_fn: Optional[Callable[..., Iterator[dict]]] = None, sleep_on_threads_shutdown: bool = False, + generate_num_examples_fn: Optional[Callable[..., Iterator[Optional[int]]]] = None, ): super().__init__() self.generate_tables_fn = generate_tables_fn self.kwargs = kwargs + # One exact count per shard, in _split_gen_kwargs(..., self.num_shards) + # order and granularity. None means that a source must be scanned. + self.generate_num_examples_fn = generate_num_examples_fn # for resharding self.generate_more_kwargs_fn = generate_more_kwargs_fn @@ -462,13 +466,22 @@ def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "A shard_indices = self.split_shard_indices_by_worker(num_shards, index, contiguous=contiguous) requested_gen_kwargs = _merge_gen_kwargs([gen_kwargs_list[i] for i in shard_indices]) return ArrowExamplesIterable( - self.generate_tables_fn, requested_gen_kwargs, self.generate_more_kwargs_fn, self.sleep_on_threads_shutdown + self.generate_tables_fn, + requested_gen_kwargs, + self.generate_more_kwargs_fn, + self.sleep_on_threads_shutdown, + self.generate_num_examples_fn, ) def reshard_data_sources(self) -> "ArrowExamplesIterable": """Split shars into more shards if possible.""" if not self.generate_more_kwargs_fn: - return ArrowExamplesIterable(self.generate_tables_fn, self.kwargs, self.generate_more_kwargs_fn) + return ArrowExamplesIterable( + self.generate_tables_fn, + self.kwargs, + self.generate_more_kwargs_fn, + generate_num_examples_fn=self.generate_num_examples_fn, + ) gen_kwargs_list = _split_gen_kwargs(self.kwargs, max_num_jobs=self.num_shards) new_gen_kwargs = _merge_gen_kwargs( [ @@ -478,7 +491,11 @@ def reshard_data_sources(self) -> "ArrowExamplesIterable": ] ) return ArrowExamplesIterable( - self.generate_tables_fn, new_gen_kwargs, self.generate_more_kwargs_fn, self.sleep_on_threads_shutdown + self.generate_tables_fn, + new_gen_kwargs, + self.generate_more_kwargs_fn, + self.sleep_on_threads_shutdown, + self.generate_num_examples_fn, ) @property @@ -2179,9 +2196,48 @@ def _init_state_dict(self) -> dict: } return self._state_dict + def _skip_complete_sources(self, skipped): + ex_iterable = self.ex_iterable + # Subclasses and wrappers may change which rows a source yields. + if ( + skipped + or self.n <= 0 + or type(ex_iterable) is not ArrowExamplesIterable + or ex_iterable.generate_num_examples_fn is None + or ( + ex_iterable._state_dict + and (ex_iterable._state_dict["shard_idx"] or ex_iterable._state_dict["shard_example_idx"]) + ) + ): + return ex_iterable, skipped + if ex_iterable._state_dict is None: + # Direct iteration without checkpoint tracking must remain repeatable. + assert self._state_dict is None + ex_iterable = copy(ex_iterable) + ex_iterable._init_state_dict() + counts = None + try: + counts = ex_iterable.generate_num_examples_fn(**ex_iterable.kwargs) + for num_examples in counts: + if num_examples is None or num_examples < 0 or skipped + num_examples > self.n: + break + skipped += num_examples + ex_iterable._state_dict["shard_idx"] += 1 + if self._state_dict: + self._state_dict["skipped"] = skipped + if skipped == self.n: + break + except NotImplementedError: + pass # A builder may not support counting this source/configuration. + finally: + if hasattr(counts, "close"): + counts.close() + return ex_iterable, skipped + def __iter__(self): skipped = self._state_dict["skipped"] if self._state_dict else 0 - for key_example in self.ex_iterable: + ex_iterable, skipped = self._skip_complete_sources(skipped) + for key_example in ex_iterable: if skipped + 1 <= self.n: skipped += 1 if self._state_dict: @@ -2191,7 +2247,8 @@ def __iter__(self): def _iter_arrow(self): skipped = self._state_dict["skipped"] if self._state_dict else 0 - for key, pa_table in self.ex_iterable.iter_arrow(): + ex_iterable, skipped = self._skip_complete_sources(skipped) + for key, pa_table in ex_iterable.iter_arrow(): if len(pa_table) == 0: continue elif skipped + len(pa_table) <= self.n: @@ -2960,6 +3017,21 @@ def _prepare_ex_iterable_for_iteration( ) -> _BaseExamplesIterable: ex_iterable = self._ex_iterable + if self._starting_state_dict and self.epoch == self._starting_state_dict["epoch"]: + # Conservatively keep checkpoint replay on the scan path, including + # rebatching that can replay a pristine source cursor. + def without_counts(ex_iterable): + ex_iterable = copy(ex_iterable) + if isinstance(ex_iterable, ArrowExamplesIterable): + ex_iterable.generate_num_examples_fn = None + if hasattr(ex_iterable, "ex_iterable"): + ex_iterable.ex_iterable = without_counts(ex_iterable.ex_iterable) + if hasattr(ex_iterable, "ex_iterables"): + ex_iterable.ex_iterables = [without_counts(child) for child in ex_iterable.ex_iterables] + return ex_iterable + + ex_iterable = without_counts(ex_iterable) + if self.epoch: try: ex_iterable = ex_iterable.shuffle_data_sources(np.random.default_rng(self.epoch)) @@ -4059,6 +4131,9 @@ def skip(self, n: int) -> "IterableDataset": """ Create a new [`IterableDataset`] that skips the first `n` elements. + Eligible Parquet streams can skip complete sources using their metadata, + without reading or validating the skipped data pages. + Args: n (`int`): Number of elements to skip. diff --git a/src/datasets/packaged_modules/parquet/parquet.py b/src/datasets/packaged_modules/parquet/parquet.py index d05ba18a298..060d9b74d96 100644 --- a/src/datasets/packaged_modules/parquet/parquet.py +++ b/src/datasets/packaged_modules/parquet/parquet.py @@ -156,6 +156,51 @@ def _cast_table(self, pa_table: pa.Table) -> pa.Table: pa_table = table_cast(pa_table, self.info.features.arrow_schema) return pa_table + def _get_examples_iterable_for_split(self, split_generator): + ex_iterable = super()._get_examples_iterable_for_split(split_generator) + # Opt in only where footer counts describe the rows that iteration yields. + # Other countable builders may require scans or count pre-filtered rows. + if ( + self.config.filters is None + and self.config.features is None + and self.config.columns is None + and self.config.on_bad_files == "error" + and (self.config.batch_size is None or self.config.batch_size > 0) + and not ( + self.config.fragment_scan_options is not None + and self.config.fragment_scan_options.page_checksum_verification + ) + ): + ex_iterable.generate_num_examples_fn = self._generate_num_examples + return ex_iterable + + def _generate_num_examples(self, files, row_groups_list): + """Yield exact source counts, or raise if a source requires the scan path.""" + previous_file, metadata = None, None + for file, row_groups in zip(files, row_groups_list): + try: + # Resharded sources can share a footer. Reuse it while visiting + # adjacent groups of a file, without caching across iterations. + if metadata is None or file != previous_file: + with open(file, "rb") as f: + metadata = pq.read_metadata(f) + previous_file = file + selected = list(range(metadata.num_row_groups)) if row_groups is None else row_groups + # Arrow normalizes duplicate/out-of-order selections. Keep these + # on its scan path rather than assuming the supplied order. + if list(selected) != sorted(set(selected)): + raise NotImplementedError + # Preserve the existing empty-first-group error until #8742 is fixed. + if selected and metadata.row_group(selected[0]).num_rows == 0 and self.config.batch_size is None: + raise NotImplementedError + # A different schema may need casting (or raise) during iteration. + if datasets.Features.from_arrow_schema(metadata.schema.to_arrow_schema()) != self.info.features: + raise NotImplementedError + yield sum(metadata.row_group(i).num_rows for i in selected) + except (OSError, pa.ArrowInvalid, ValueError, TypeError, KeyError, IndexError) as e: + # Let normal iteration handle unsupported metadata and bad files. + raise NotImplementedError from e + def _generate_shards(self, files, row_groups_list): if not row_groups_list: yield from files diff --git a/tests/packaged_modules/test_parquet.py b/tests/packaged_modules/test_parquet.py index e0206b5b53f..70089417851 100644 --- a/tests/packaged_modules/test_parquet.py +++ b/tests/packaged_modules/test_parquet.py @@ -1,10 +1,366 @@ +from copy import deepcopy + +import pyarrow as pa +import pyarrow.dataset as pds +import pyarrow.parquet as pq import pytest -from datasets import load_dataset +from datasets import Features, Value, concatenate_datasets, load_dataset from datasets.builder import InvalidConfigName from datasets.data_files import DataFilesList +from datasets.distributed import split_dataset_by_node +from datasets.iterable_dataset import SkipExamplesIterable from datasets.packaged_modules.parquet.parquet import ParquetConfig +from ..utils import require_torch + + +@pytest.fixture +def parquet_skip_paths(tmp_path): + files = [] + for i, values in enumerate([range(6), range(6, 11), range(11, 15)]): + path = tmp_path / f"{i}.parquet" + pq.write_table(pa.table({"id": values}), path, row_group_size=2) + files.append(str(path)) + return files + + +@pytest.fixture( + params=[(False, None), (False, "arrow"), (True, None), (True, "arrow")], + ids=["files-python", "files-arrow", "row-groups-python", "row-groups-arrow"], +) +def parquet_stream(request, parquet_skip_paths): + # Unequal file lengths and a partial final row group exercise both source granularities. + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + reshard, formatting = request.param + if reshard: + dataset = dataset.reshard() + return dataset.with_format(formatting) + + +def _parquet_stream_ids(examples): + return [ + row["id"] + for example in examples + for row in (example.to_pylist() if isinstance(example, pa.Table) else [example]) + ] + + +def test_parquet_stream_skip_boundaries(parquet_stream): + # Before, at, and after row-group/file boundaries; the last two skips exhaust the stream. + for n in [0, 1, 2, 3, 5, 6, 7, 10, 11, 12, 14, 15, 16]: + assert _parquet_stream_ids(parquet_stream.skip(n)) == list(range(n, 15)), n + + +@pytest.mark.parametrize("consumed", [1, 5], ids=["first-retained-row", "across-file-boundary"]) +def test_parquet_stream_skip_resume(parquet_stream, consumed): + # Skip a whole file and part of a row group before saving progress. + dataset = parquet_stream.skip(7) + iterator = iter(dataset) + prefix = [next(iterator) for _ in range(consumed)] + state = deepcopy(dataset.state_dict()) + remaining = _parquet_stream_ids(iterator) + + restored = parquet_stream.skip(7) + restored.load_state_dict(state) + assert _parquet_stream_ids(prefix) == list(range(7, 7 + consumed)) + assert remaining == list(range(7 + consumed, 15)) + assert _parquet_stream_ids(restored) == remaining + + +@pytest.mark.parametrize("consumed_batches", [0, 1, 3]) +def test_parquet_stream_skip_resume_batches(parquet_stream, consumed_batches): + def batch_ids(batches): + return [ + value + for batch in batches + for value in (batch["id"].to_pylist() if isinstance(batch, pa.Table) else batch["id"]) + ] + + dataset = parquet_stream.skip(7) + iterator = dataset.iter(batch_size=3) + prefix = [next(iterator) for _ in range(consumed_batches)] + # Force exhaustion for the checkpoint after the final partial batch. + if consumed_batches == 3: + assert next(iterator, None) is None + state = deepcopy(dataset.state_dict()) + tail = batch_ids(iterator) + restored = parquet_stream.skip(7) + restored.load_state_dict(state) + assert batch_ids(prefix) + tail == list(range(7, 15)) + assert batch_ids(restored.iter(batch_size=3)) == tail + + +def test_parquet_stream_skip_repeated_epochs(parquet_stream): + dataset = parquet_stream.skip(7) + for epoch in [0, 1, 1, 2, 0]: + dataset.set_epoch(epoch) + assert _parquet_stream_ids(dataset) == list(range(7, 15)), epoch + + +@pytest.mark.parametrize("n, expected_decoded", [(7, 2), (15, 0), (16, 0)]) +def test_parquet_stream_skip_does_not_decode_complete_sources(parquet_stream, monkeypatch, n, expected_decoded): + source = parquet_stream._ex_iterable + generate_tables = source.generate_tables_fn + decoded = [] + + def counted_tables(**kwargs): + for key, table in generate_tables(**kwargs): + decoded.append(len(table)) + yield key, table + + monkeypatch.setattr(source, "generate_tables_fn", counted_tables) + iterator = iter(parquet_stream.skip(n)) + first = next(iterator, None) + assert _parquet_stream_ids([] if first is None else [first]) == ([n] if n < 15 else []) + assert sum(decoded) == expected_decoded + + +def test_parquet_skip_selected_row_groups(parquet_skip_paths): + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + source = dataset._ex_iterable + source.kwargs = {"files": parquet_skip_paths, "row_groups_list": [(1, 2), (1,), ()]} + assert list(source.generate_num_examples_fn(**source.kwargs)) == [4, 2, 0] + expected = [2, 3, 4, 5, 8, 9] + for n in [0, 3, 4, 5, 6, 7]: + assert _parquet_stream_ids(dataset.skip(n)) == expected[n:] + + +@pytest.mark.parametrize("operation", ["map", "filter", "shuffle"]) +def test_parquet_skip_transformed_input_falls_back(parquet_stream, monkeypatch, operation): + def must_not_count(**kwargs): + pytest.fail("Transformed inputs must use the scan path") + + monkeypatch.setattr(parquet_stream._ex_iterable, "generate_num_examples_fn", must_not_count) + dataset = parquet_stream.with_format(None) + if operation == "map": + dataset = dataset.map(lambda row: {"id": row["id"] * 2}) + elif operation == "filter": + dataset = dataset.filter(lambda row: row["id"] % 2 == 0) + else: + dataset = dataset.shuffle(seed=42, buffer_size=4) + expected = _parquet_stream_ids(dataset) + assert _parquet_stream_ids(dataset.skip(7)) == expected[7:] + + +@pytest.mark.parametrize( + "config", + [ + {"filters": [("id", ">", 3)]}, + {"columns": ["id"]}, + {"features": Features({"id": Value("float64")})}, + {"on_bad_files": "skip"}, + {"on_bad_files": "warn"}, + {"fragment_scan_options": pds.ParquetFragmentScanOptions(page_checksum_verification=True)}, + ], +) +def test_parquet_skip_config_falls_back(parquet_skip_paths, monkeypatch, config): + from datasets.packaged_modules.parquet.parquet import Parquet + + def must_not_count(*args, **kwargs): + pytest.fail("Unsupported configurations must use the scan path") + + monkeypatch.setattr(Parquet, "_generate_num_examples", must_not_count) + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True, **config) + assert _parquet_stream_ids(dataset.skip(7)) == _parquet_stream_ids(dataset)[7:] + + +def test_parquet_skip_restored_state_falls_back(parquet_stream, monkeypatch): + dataset = parquet_stream.skip(7) + iterator = iter(dataset) + next(iterator) + state = deepcopy(dataset.state_dict()) + + def must_not_count(**kwargs): + pytest.fail("Checkpoint replay must use the scan path") + + monkeypatch.setattr(parquet_stream._ex_iterable, "generate_num_examples_fn", must_not_count) + restored = parquet_stream.skip(7) + restored.load_state_dict(state) + assert _parquet_stream_ids(restored) == list(range(8, 15)) + + +def test_parquet_skip_progressed_source_falls_back(parquet_stream, monkeypatch): + source = parquet_stream._ex_iterable + source._init_state_dict() + iterator = iter(source) + assert next(iterator)[1]["id"] == 0 + iterator.close() + + def must_not_count(**kwargs): + pytest.fail("An already progressed source must use the scan path") + + monkeypatch.setattr(source, "generate_num_examples_fn", must_not_count) + assert [row["id"] for _, row in SkipExamplesIterable(source, 7)] == list(range(8, 15)) + + +def test_parquet_skip_concatenated_restore_falls_back(parquet_stream, monkeypatch): + dataset = concatenate_datasets([parquet_stream.skip(7), parquet_stream.skip(7)]) + iterator = iter(dataset) + next(iterator) + state = deepcopy(dataset.state_dict()) + + def must_not_count(**kwargs): + pytest.fail("Restored children must use the scan path, including an unvisited child") + + monkeypatch.setattr(parquet_stream._ex_iterable, "generate_num_examples_fn", must_not_count) + restored = concatenate_datasets([parquet_stream.skip(7), parquet_stream.skip(7)]) + restored.load_state_dict(state) + assert _parquet_stream_ids(restored) == list(range(8, 15)) + list(range(7, 15)) + + +@pytest.mark.parametrize("num_shards, contiguous", [(2, True), (2, False), (5, False)]) +def test_parquet_skip_worker_sharding_matches_scan(parquet_stream, num_shards, contiguous): + fast = SkipExamplesIterable(parquet_stream._ex_iterable, 7) + slow_source = deepcopy(parquet_stream._ex_iterable) + slow_source.generate_num_examples_fn = None + slow = SkipExamplesIterable(slow_source, 7) + for index in range(min(num_shards, fast.num_shards)): + fast_worker = fast.shard_data_sources(num_shards, index, contiguous) + slow_worker = slow.shard_data_sources(num_shards, index, contiguous) + assert list(fast_worker) == list(slow_worker) + + +@require_torch +@pytest.mark.parametrize("num_workers, reshard", [(2, False), (5, False), (2, True)]) +def test_parquet_skip_dataloader_matches_scan(parquet_skip_paths, num_workers, reshard): + from torch.utils.data import DataLoader + + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + if reshard: + dataset = dataset.reshard() + slow = deepcopy(dataset) + slow._ex_iterable.generate_num_examples_fn = None + actual = sorted(_parquet_stream_ids(DataLoader(dataset.skip(7), batch_size=None, num_workers=num_workers))) + expected = sorted(_parquet_stream_ids(DataLoader(slow.skip(7), batch_size=None, num_workers=num_workers))) + assert actual == expected + + +@pytest.mark.parametrize("world_size, skip_first", [(2, True), (3, True), (2, False), (3, False)]) +def test_parquet_skip_distributed_matches_scan(parquet_stream, world_size, skip_first): + slow = deepcopy(parquet_stream) + slow._ex_iterable.generate_num_examples_fn = None + for rank in range(world_size): + + def prepare(dataset): + if skip_first: + dataset = dataset.skip(7) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return dataset if skip_first else dataset.skip(7) + + assert _parquet_stream_ids(prepare(parquet_stream)) == _parquet_stream_ids(prepare(slow)) + + +def test_parquet_skip_unknown_metadata(parquet_stream, monkeypatch): + original = pq.read_metadata + calls = [] + + def unavailable(file, *args, **kwargs): + calls.append(file.name) + if file.name.endswith("1.parquet"): + raise OSError("metadata unavailable") + return original(file, *args, **kwargs) + + monkeypatch.setattr(pq, "read_metadata", unavailable) + assert _parquet_stream_ids(parquet_stream.skip(12)) == [12, 13, 14] + assert calls[-1].endswith("1.parquet") + + +def test_parquet_skip_reuses_footer_after_resharding(parquet_skip_paths, monkeypatch): + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True).reshard() + original = pq.read_metadata + calls = [] + + def counted(file, *args, **kwargs): + calls.append(file.name) + return original(file, *args, **kwargs) + + monkeypatch.setattr(pq, "read_metadata", counted) + assert _parquet_stream_ids(dataset.skip(16)) == [] + assert calls == parquet_skip_paths + assert _parquet_stream_ids(dataset.skip(0)) == list(range(15)) + assert calls == parquet_skip_paths + + +@pytest.mark.parametrize("groups", [[], [0], [0, 2], [2, 0, 2]]) +def test_parquet_skip_empty_sources_preserves_baseline(tmp_path, groups): + path = tmp_path / "empty.parquet" + with pq.ParquetWriter(path, pa.schema([("id", pa.int64())])) as writer: + for size in groups: + writer.write_table(pa.table({"id": pa.array(range(size), type=pa.int64())})) + dataset = load_dataset("parquet", data_files=str(path), split="train", streaming=True) + if groups and groups[0] == 0: + # Existing main behavior, owned by #8742; metadata skipping must not hide it. + with pytest.raises(pa.ArrowInvalid, match="BatchSize must be greater than 0"): + list(dataset.skip(10)) + else: + assert list(dataset.skip(10)) == [] + + +def test_parquet_skip_required_fields(parquet_skip_paths, monkeypatch): + schema = pa.schema([pa.field("id", pa.int64(), nullable=False)]) + pq.write_table(pa.Table.from_arrays([pa.array(range(6))], schema=schema), parquet_skip_paths[0], row_group_size=2) + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + source = dataset._ex_iterable + assert sum(source.generate_num_examples_fn(**source.kwargs)) == 15 + original = source.generate_tables_fn + scanned = [] + + def tables(**kwargs): + scanned.extend(kwargs["files"]) + yield from original(**kwargs) + + monkeypatch.setattr(source, "generate_tables_fn", tables) + assert _parquet_stream_ids(dataset.skip(7)) == list(range(7, 15)) + assert parquet_skip_paths[0] not in scanned + + +def test_parquet_skip_schema_change_falls_back(parquet_skip_paths, monkeypatch): + pq.write_table(pa.table({"id": pa.array(range(6, 11), type=pa.int32())}), parquet_skip_paths[1], row_group_size=2) + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + source = dataset._ex_iterable + counts = source.generate_num_examples_fn(**source.kwargs) + assert next(counts) == 6 + with pytest.raises(NotImplementedError): + next(counts) + original = source.generate_tables_fn + decoded = [] + + def tables(**kwargs): + for key, table in original(**kwargs): + decoded.append(len(table)) + yield key, table + + monkeypatch.setattr(source, "generate_tables_fn", tables) + assert next(iter(dataset.skip(12)))["id"] == 12 + assert sum(decoded) == 7 # Scan the differently typed source, then the retained batch. + + +def test_parquet_skip_bad_footer_still_raises(parquet_skip_paths): + with open(parquet_skip_paths[1], "wb") as file: + file.write(b"not a parquet footer") + dataset = load_dataset("parquet", data_files=parquet_skip_paths, split="train", streaming=True) + with pytest.raises(pa.ArrowInvalid): + list(dataset.skip(15)) + + +@pytest.mark.parametrize("selection", [(1, 0), (1, 1)]) +def test_parquet_skip_noncanonical_row_groups_falls_back(tmp_path, selection): + path = tmp_path / "groups.parquet" + with pq.ParquetWriter(path, pa.schema([("id", pa.int64())])) as writer: + writer.write_table(pa.table({"id": pa.array([], type=pa.int64())})) + writer.write_table(pa.table({"id": [1, 2]})) + dataset = load_dataset("parquet", data_files=str(path), split="train", streaming=True) + source = dataset._ex_iterable + source.kwargs["row_groups_list"] = [selection] + with pytest.raises(NotImplementedError): + list(source.generate_num_examples_fn(**source.kwargs)) + if 0 in selection: + with pytest.raises(pa.ArrowInvalid, match="BatchSize must be greater than 0"): + list(dataset.skip(2)) + else: + assert list(dataset.skip(2)) == [] + def test_config_raises_when_invalid_name() -> None: with pytest.raises(InvalidConfigName, match="Bad characters"): diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py index 3d14a450463..90ef5d7350e 100644 --- a/tests/test_iterable_dataset.py +++ b/tests/test_iterable_dataset.py @@ -1650,6 +1650,49 @@ def test_skip_arrow_examples_iterable(count): assert_load_state_dict_resumes_arrow_iteration(skip_ex_iterable) +@pytest.mark.parametrize("arrow", [False, True]) +@pytest.mark.parametrize("counting", ["known", "unknown", "unavailable", "missing"]) +def test_skip_arrow_complete_sources(arrow, counting): + kwargs = {"filepaths": ["a", "b", "c"], "n": 5, "batch_size": 2} + generated = [] + + def tables(**kwargs): + generated.extend(kwargs["filepaths"]) + yield from generate_tables_fn(**kwargs) + + def counts(**kwargs): + if counting == "unavailable": + raise NotImplementedError + return iter([5, None, 5] if counting == "unknown" else [5, 5, 5]) + + source = ArrowExamplesIterable(tables, kwargs, generate_num_examples_fn=None if counting == "missing" else counts) + skipped = SkipExamplesIterable(source, 11) + + def collect(ex_iterable): + if arrow: + return [row for _, table in ex_iterable.iter_arrow() for row in table.to_pylist()] + return [row for _, row in ex_iterable] + + expected = [row for _, table in generate_tables_fn(**kwargs) for row in table.to_pylist()][11:] + assert collect(skipped) == expected + assert ( + generated + == {"known": ["c"], "unknown": ["b", "c"], "unavailable": ["a", "b", "c"], "missing": ["a", "b", "c"]}[ + counting + ] + ) + # Iteration without explicit checkpoint initialization must remain reusable. + assert collect(skipped) == expected + skipped._init_state_dict() + assert collect(skipped) == expected + assert skipped.state_dict()["skipped"] == 11 + assert skipped.state_dict()["examples_iterable"]["shard_idx"] == 3 + assert_arrow = ( + assert_load_state_dict_resumes_arrow_iteration if arrow else assert_load_state_dict_resumes_iteration + ) + assert_arrow(skipped) + + def test_take_arrow_examples_iterable(): total, count = 10, 2 base_ex_iterable = ArrowExamplesIterable(generate_tables_fn, {"n": total})