Skip complete Parquet sources using metadata counts

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.

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