



'''simple docstring'''






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     snake_case_					(     lowerCAmelCase_       ,   lowerCAmelCase_		)->		Dict:





       '''simple docstring'''


       _UpperCAmelCase  :						Optional[Any]			     = []
       for part_id in partition_order:
              _UpperCAmelCase  :						Tuple			     = df.where(F'''SPARK_PARTITION_ID() = {part_id}'''		).collect()
              for row_idx, row in enumerate(lowerCAmelCase_		):
                     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     snake_case_					(     )->		int:





       '''simple docstring'''


       _UpperCAmelCase  :						List[Any]			     = pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
       _UpperCAmelCase  :						Dict			     = spark.range(100		).repartition(1		)
       _UpperCAmelCase  :						Optional[int]			     = Spark(lowerCAmelCase_		)
       # 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     snake_case_					(     )->		str:





       '''simple docstring'''


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

       for i, (row_id, row_dict) in enumerate(generate_fn()		):
              _UpperCAmelCase     ,_UpperCAmelCase  :						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     snake_case_					(     )->		List[Any]:





       '''simple docstring'''


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




@require_not_windows
@require_dill_gt_0_3_2
def     snake_case_					(     )->		List[str]:





       '''simple docstring'''


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

              _UpperCAmelCase  :						Union[str, Any]			     = SparkExamplesIterable(lowerCAmelCase_		).shuffle_data_sources(lowerCAmelCase_		)
              assert shuffled_it.n_shards == 3
              for i, (row_id, row_dict) in enumerate(lowerCAmelCase_		):
                     _UpperCAmelCase     ,_UpperCAmelCase  :						List[str]			     = 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     snake_case_					(     )->		int:





       '''simple docstring'''


       _UpperCAmelCase  :						Optional[int]			     = pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
       _UpperCAmelCase  :						Tuple			     = spark.range(20		).repartition(4		)

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

       # Partitions 1 and 3
       _UpperCAmelCase  :						Dict			     = SparkExamplesIterable(lowerCAmelCase_		).shard_data_sources(worker_id=1       ,   num_workers=2		)
       assert shard_it_a.n_shards == 2
       _UpperCAmelCase  :						str			     = _get_expected_row_ids_and_row_dicts_for_partition_order(lowerCAmelCase_       ,   [1, 3]		)
       for i, (row_id, row_dict) in enumerate(lowerCAmelCase_		):
              _UpperCAmelCase     ,_UpperCAmelCase  :						Any			     = 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     snake_case_					(     )->		Union[str, Any]:





       '''simple docstring'''


       _UpperCAmelCase  :						List[str]			     = pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
       _UpperCAmelCase  :						Any			     = spark.range(100		).repartition(1		)
       _UpperCAmelCase  :						Optional[int]			     = Spark(lowerCAmelCase_		)
       # 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








'''simple docstring'''






from numpy import exp, pi, sqrt




def     snake_case_					(     lowerCAmelCase_       ,   lowerCAmelCase_ = 0.0       ,   lowerCAmelCase_ = 1.0		)->		int:





       '''simple docstring'''


       return 1 / sqrt(2 * pi * sigma**2		) * exp(-((x - mu) ** 2) / (2 * sigma**2)		)


if __name__ == "__main__":
   import doctest

   doctest.testmod()



