class SparkExamplesIterable(_BaseExamplesIterable):
    def __init__(
        self,
        df: "pyspark.sql.DataFrame",
        partition_order=None,
    ):
        super().__init__()
        self.df = df
        self.partition_order = partition_order or range(self.df.rdd.getNumPartitions())

    def _init_state_dict(self) -> dict:
        self._state_dict = {"partition_idx": 0, "partition_example_idx": 0}
        return self._state_dict

    @experimental
    def load_state_dict(self, state_dict: dict) -> dict:
        return super().load_state_dict(state_dict)

    def __iter__(self):
        yield from _generate_iterable_examples(self.df, self.partition_order, self._state_dict)

    def shuffle_data_sources(self, generator: np.random.Generator) -> "SparkExamplesIterable":
        partition_order = list(range(self.df.rdd.getNumPartitions()))
        generator.shuffle(partition_order)
        return SparkExamplesIterable(self.df, partition_order=partition_order)

    def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "SparkExamplesIterable":
        partition_order = self.split_shard_indices_by_worker(num_shards=num_shards, index=index, contiguous=contiguous)
        return SparkExamplesIterable(self.df, partition_order=partition_order)

    @property
    def num_shards(self) -> int:
        return len(self.partition_order)
/Users/nielsrogge/Documents/python_projecten/datasets/src/datasets/packaged_modules/spark/spark.py