Repository: huggingface/datasets
File: src/datasets/iterable_dataset.py (Python)

Code before:

        if sys.platform == "win32":
            raise EnvironmentError("IterableDataset.from_spark is not currently supported on Windows")

        return SparkDatasetReader(
            df,
            split=split,
            features=features,
            streaming=True,
            **kwargs,
        ).read()

    @staticmethod
    def from_file(filename: str, in_memory: bool = False) -> "IterableDataset":
        """Instantiate a IterableDataset from Arrow table at filename.

        Args:
            filename (`str`):
                File name of the dataset.
            in_memory (`bool`, defaults to `False`):
                Whether to copy the data in-memory.

        Returns:
            [`IterableDataset`]
        """
        pa_table = ArrowReader.read_table(filename, in_memory=in_memory)
        inferred_features = Features.from_arrow_schema(pa_table.schema)
        ex_iterable = ExamplesIterable(Dataset._generate_examples_from_table, kwargs={"arrow_table": pa_table})
        return IterableDataset(ex_iterable=ex_iterable, info=DatasetInfo(features=inferred_features))

    def with_format(
        self,
        type: Optional[str] = None,
    ) -> "IterableDataset":
        """
        Return a dataset with the specified format.
        This method only supports the "torch" format for now.

        Args:

            type (`str`, optional, default None): if set to "torch", the returned dataset
                will be a subclass of torch.utils.data.IterableDataset to be used in a DataLoader
        """
        type = get_format_type_from_alias(type)
        # TODO(QL): add examples formatting to get tensors when using the "torch" format
        # TODO(QL): add format_kwargs
        # TODO(QL): add format_columns and return_all_columns
        # TODO(QL): add pandas, numpy and tf formats
        return IterableDataset(
            ex_iterable=self._ex_iterable,
            info=self._info.copy(),

Review comment (performance):
This would read the entire file metadata, which can take some time, in particular if you have thousands of files (which is going to be the case for people interested in https://github.com/huggingface/datasets/issues/5481)

I think it should be possible to have an arrow table generator that yields tables from the file without reading all its metadata.

Maybe using `pyarrow.ipc.RecordBatchStreamReader` and then `read_next_batch` ?

Code after:
        import torch.utils.data

        if torch.utils.data.IterableDataset not in cls.__bases__:
            cls.__bases__ += (torch.utils.data.IterableDataset,)


class IterableDataset(DatasetInfoMixin):
    """A Dataset backed by an iterable."""

    def __init__(
        self,
        ex_iterable: _BaseExamplesIterable,
        info: Optional[DatasetInfo] = None,
        split: Optional[NamedSplit] = None,
        format_type: Optional[str] = None,
        shuffling: Optional[ShufflingConfig] = None,
        distributed: Optional[DistributedConfig] = None,
        token_per_repo_id: Optional[Dict[str, Union[str, bool, None]]] = None,
    ):
        if distributed and distributed.world_size > 1 and shuffling and shuffling._original_seed is None:
            raise RuntimeError(
                "The dataset doesn't have a fixed random seed across nodes to shuffle and split the list of dataset shards by node. "
                "Please pass e.g. `seed=42` in `.shuffle()` to make all the nodes use the same seed. "
            )

        info = info.copy() if info is not None else DatasetInfo()
        DatasetInfoMixin.__init__(self, info=info, split=split)

        self._ex_iterable = ex_iterable
        self._format_type = format_type
        self._shuffling = shuffling
        self._distributed = distributed
        self._epoch = 0
        self._token_per_repo_id: Dict[str, Union[str, bool, None]] = token_per_repo_id or {}
        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)

    def __getstate__(self):
        return self.__dict__

    def __setstate__(self, d):
        self.__dict__ = d
        # Re-add torch iterable dataset as a parent class, since dynamically added parent classes are not kept when pickling
        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)

    def _head(self, n=5):
        return _examples_to_batch(list(self.take(n)))

    def _effective_generator(self):
        if self._shuffling and self._epoch == 0:
            return self._shuffling.generator
        elif self._shuffling: