Method _create_iterator

python/ray/data/iterator.py:245–307

()

        batch_format = _apply_batch_format(batch_format)

        def _create_iterator() -> Iterator[DataBatch]:
            time_start = time.perf_counter()

            # Iterate through the dataset from the start each time
            # _iterator_gen is called.
            # This allows multiple iterations of the dataset without
            # needing to explicitly call `iter_batches()` multiple times.
            (
                ref_bundles_iterator,
                stats,
                blocks_owned_by_consumer,
                executor,
            ) = self._to_ref_bundle_iterator()

            dataset_tag = self._get_dataset_tag()

            # Create a callback to report prefetched bytes to the executor's
            # resource manager.
            def make_prefetch_callback(exec):
                def callback(num_bytes: int) -> None:
                    exec.set_external_consumer_bytes(num_bytes)

                return callback

            prefetch_bytes_callback = (
                make_prefetch_callback(executor) if executor is not None else None
            )
            if prefetch_bytes_callback is not None:
                # Register the external consumer with the executor's resource manager.
                prefetch_bytes_callback(0)

            batch_iterator = self._create_batch_iterator(
                ref_bundles_iterator,
                stats=stats,
                dataset_tag=dataset_tag,
                clear_block_after_read=blocks_owned_by_consumer,
                batch_size=batch_size,
                batch_format=batch_format,
                drop_last=drop_last,
                collate_fn=_collate_fn,
                finalize_fn=_finalize_fn,
                shuffle_buffer_min_size=local_shuffle_buffer_size,
                shuffle_seed=local_shuffle_seed,
                prefetch_batches=prefetch_batches,
                prefetch_bytes_callback=prefetch_bytes_callback,
                preserve_order=self.get_context().execution_options.preserve_order,
            )

            if stats:
                stats.iter_initialize_s.add(time.perf_counter() - time_start)

            try:
                yield from batch_iterator
                if stats:
                    stats.iter_total_s.add(time.perf_counter() - time_start)
            finally:
                # On early exit (e.g. ``break`` in the for-loop), the inner
                # ``_ClosingIterator`` would only shut down the executor via