


from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)





def   _UpperCAmelCase       (     UpperCAmelCase      :			List[str]						,	UpperCAmelCase      :			List[Any]      ):
      """simple docstring"""




      __lowerCamelCase     :      Tuple       						=						[]
      for part_id in partition_order:
            __lowerCamelCase     :      List[Any]       						=						df.where(f"""SPARK_PARTITION_ID() = {part_id}"""      ).collect()
            for row_idx, row in enumerate(UpperCAmelCase      ):
                  expected_row_ids_and_row_dicts.append((f"""{part_id}_{row_idx}""", row.asDict())      )
      return expected_row_ids_and_row_dicts





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      int       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      Dict       						=						spark.range(100      ).repartition(1      )
      __lowerCamelCase     :      Dict       						=						Spark(UpperCAmelCase      )
      # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
      # that each partition can hold 2 rows.
      spark_builder._repartition_df_if_needed(max_shard_size=16      )
      # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
      assert spark_builder.df.rdd.getNumPartitions() == 50





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      Optional[Any]       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      Union[str, Any]       						=						spark.range(10      ).repartition(2      )
      __lowerCamelCase     :      Any       						=						[1, 0]
      __lowerCamelCase     :      Any       						=						_generate_iterable_examples(UpperCAmelCase						,	UpperCAmelCase      )  # Reverse the partitions.
      __lowerCamelCase     :      Any       						=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase						,	UpperCAmelCase      )

      for i, (row_id, row_dict) in enumerate(generate_fn()      ):
            __lowerCamelCase						, __lowerCamelCase     :      Optional[Any]       						=						expected_row_ids_and_row_dicts[i]
            assert row_id == expected_row_id
            assert row_dict == expected_row_dict





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      Union[str, Any]       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      Optional[int]       						=						spark.range(10      ).repartition(1      )
      __lowerCamelCase     :      Optional[Any]       						=						SparkExamplesIterable(UpperCAmelCase      )
      assert it.n_shards == 1
      for i, (row_id, row_dict) in enumerate(UpperCAmelCase      ):
            assert row_id == f"""0_{i}"""
            assert row_dict == {"id": i}





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      List[str]       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      str       						=						spark.range(30      ).repartition(3      )
      # Mock the generator so that shuffle reverses the partition indices.
      with patch("""numpy.random.Generator"""      ) as generator_mock:
            __lowerCamelCase     :      str       						=						lambda UpperCAmelCase      : x.reverse()
            __lowerCamelCase     :      Optional[int]       						=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase						,	[2, 1, 0]      )

            __lowerCamelCase     :      List[Any]       						=						SparkExamplesIterable(UpperCAmelCase      ).shuffle_data_sources(UpperCAmelCase      )
            assert shuffled_it.n_shards == 3
            for i, (row_id, row_dict) in enumerate(UpperCAmelCase      ):
                  __lowerCamelCase						, __lowerCamelCase     :      Tuple       						=						expected_row_ids_and_row_dicts[i]
                  assert row_id == expected_row_id
                  assert row_dict == expected_row_dict





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      int       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      List[Any]       						=						spark.range(20      ).repartition(4      )

      # Partitions 0 and 2
      __lowerCamelCase     :      List[str]       						=						SparkExamplesIterable(UpperCAmelCase      ).shard_data_sources(worker_id=0						,	num_workers=2      )
      assert shard_it_a.n_shards == 2
      __lowerCamelCase     :      List[Any]       						=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase						,	[0, 2]      )
      for i, (row_id, row_dict) in enumerate(UpperCAmelCase      ):
            __lowerCamelCase						, __lowerCamelCase     :      Tuple       						=						expected_row_ids_and_row_dicts_a[i]
            assert row_id == expected_row_id
            assert row_dict == expected_row_dict

      # Partitions 1 and 3
      __lowerCamelCase     :      Optional[int]       						=						SparkExamplesIterable(UpperCAmelCase      ).shard_data_sources(worker_id=1						,	num_workers=2      )
      assert shard_it_a.n_shards == 2
      __lowerCamelCase     :      List[Any]       						=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase						,	[1, 3]      )
      for i, (row_id, row_dict) in enumerate(UpperCAmelCase      ):
            __lowerCamelCase						, __lowerCamelCase     :      Dict       						=						expected_row_ids_and_row_dicts_a[i]
            assert row_id == expected_row_id
            assert row_dict == expected_row_dict





@require_not_windows
@require_dill_gt_0_3_2
def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      int       						=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
      __lowerCamelCase     :      int       						=						spark.range(100      ).repartition(1      )
      __lowerCamelCase     :      List[str]       						=						Spark(UpperCAmelCase      )
      # Choose a small max_shard_size for maximum partitioning.
      spark_builder._repartition_df_if_needed(max_shard_size=1      )
      # The new number of partitions should not be greater than the number of rows.
      assert spark_builder.df.rdd.getNumPartitions() == 100







from typing import Any





def   _UpperCAmelCase       (     UpperCAmelCase      :			list      ):
      """simple docstring"""




      if not input_list:
            return []
      __lowerCamelCase     :      int       						=						[input_list.count(UpperCAmelCase      ) for value in input_list]
      __lowerCamelCase     :      Tuple       						=						max(UpperCAmelCase      )  # Gets the maximum count in the input list.
      # Gets values of modes
      return sorted({input_list[i] for i, value in enumerate(UpperCAmelCase      ) if value == y}      )


if __name__ == "__main__":
    import doctest

    doctest.testmod()



