fix tests

Author: lhoestq

src/datasets/packaged_modules/generator/generator.py

@@ -32,6 +32,6 @@ def _split_generators(self, dl_manager):
 
     def _generate_examples(self, **gen_kwargs):
         num_shards = _number_of_shards_in_gen_kwargs(gen_kwargs)
-        for shard_idx, shard_gen_kwargs in _split_gen_kwargs(gen_kwargs, max_num_jobs=num_shards):
+        for shard_idx, shard_gen_kwargs in enumerate(_split_gen_kwargs(gen_kwargs, max_num_jobs=num_shards)):
             for sample_idx, sample in enumerate(self.config.generator(**shard_gen_kwargs)):
                 yield (shard_idx, sample_idx), sample

tests/packaged_modules/test_spark.py

@@ -25,7 +25,7 @@ def _get_expected_row_ids_and_row_dicts_for_partition_order(df, partition_order)
     for part_id in partition_order:
         partition = df.where(f"SPARK_PARTITION_ID() = {part_id}").collect()
         for row_idx, row in enumerate(partition):
-            expected_row_ids_and_row_dicts.append((f"{part_id}_{row_idx}", row.asDict()))
+            expected_row_ids_and_row_dicts.append(((part_id, row_idx), row.asDict()))
     return expected_row_ids_and_row_dicts