



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 (       __snake_case					,  __snake_case    )     ->     Optional[int]:






      '''simple docstring'''







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






      '''simple docstring'''







      UpperCAmelCase_  :	Dict				      =		pyspark.sql.SparkSession.builder.master("local[*]"    ).appName("pyspark"    ).getOrCreate()
      UpperCAmelCase_  :	str				      =		spark.range(1_0_0    ).repartition(1    )
      UpperCAmelCase_  :	Optional[Any]				      =		Spark(__snake_case    )
      # 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=1_6    )
      # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
      assert spark_builder.df.rdd.getNumPartitions() == 5_0







@require_not_windows
@require_dill_gt_0_3_2
def    _snake_case (       )     ->     Tuple:






      '''simple docstring'''







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

      for i, (row_id, row_dict) in enumerate(generate_fn()    ):
            UpperCAmelCase_      ,      UpperCAmelCase_  :	int				      =		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 (       )     ->     str:






      '''simple docstring'''







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







@require_not_windows
@require_dill_gt_0_3_2
def    _snake_case (       )     ->     str:






      '''simple docstring'''







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

            UpperCAmelCase_  :	Any				      =		SparkExamplesIterable(__snake_case    ).shuffle_data_sources(__snake_case    )
            assert shuffled_it.n_shards == 3
            for i, (row_id, row_dict) in enumerate(__snake_case    ):
                  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 (       )     ->     List[Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Dict				      =		pyspark.sql.SparkSession.builder.master("local[*]"    ).appName("pyspark"    ).getOrCreate()
      UpperCAmelCase_  :	int				      =		spark.range(2_0    ).repartition(4    )

      # Partitions 0 and 2
      UpperCAmelCase_  :	Optional[int]				      =		SparkExamplesIterable(__snake_case    ).shard_data_sources(worker_id=0					,  num_workers=2    )
      assert shard_it_a.n_shards == 2
      UpperCAmelCase_  :	Optional[int]				      =		_get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case					,  [0, 2]    )
      for i, (row_id, row_dict) in enumerate(__snake_case    ):
            UpperCAmelCase_      ,      UpperCAmelCase_  :	List[str]				      =		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_  :	Optional[int]				      =		SparkExamplesIterable(__snake_case    ).shard_data_sources(worker_id=1					,  num_workers=2    )
      assert shard_it_a.n_shards == 2
      UpperCAmelCase_  :	Optional[int]				      =		_get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case					,  [1, 3]    )
      for i, (row_id, row_dict) in enumerate(__snake_case    ):
            UpperCAmelCase_      ,      UpperCAmelCase_  :	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    _snake_case (       )     ->     Dict:






      '''simple docstring'''







      UpperCAmelCase_  :	Optional[int]				      =		pyspark.sql.SparkSession.builder.master("local[*]"    ).appName("pyspark"    ).getOrCreate()
      UpperCAmelCase_  :	Optional[int]				      =		spark.range(1_0_0    ).repartition(1    )
      UpperCAmelCase_  :	int				      =		Spark(__snake_case    )
      # 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() == 1_0_0





import operator







def    _snake_case (       __snake_case					,  __snake_case = False					,  __snake_case = None    )     ->     list:






      '''simple docstring'''







      UpperCAmelCase_  :	Optional[int]				      =		operator.lt if reverse else operator.gt
      UpperCAmelCase_  :	int				      =		solution or []

      if not arr:
            return solution

      UpperCAmelCase_  :	Union[str, Any]				      =		[arr.pop(0    )]
      for i, item in enumerate(__snake_case    ):
            if _operator(__snake_case					,  sublist[-1]    ):
                  sublist.append(__snake_case    )
                  arr.pop(__snake_case    )

    #  merging sublist into solution list
      if not solution:
            solution.extend(__snake_case    )
      else:
            while sublist:
                  UpperCAmelCase_  :	Optional[int]				      =		sublist.pop(0    )
                  for i, xx in enumerate(__snake_case    ):
                        if not _operator(__snake_case					,  __snake_case    ):
                              solution.insert(__snake_case					,  __snake_case    )
                              break
                  else:
                        solution.append(__snake_case    )

      strand_sort(__snake_case					,  __snake_case					,  __snake_case    )
      return solution


if __name__ == "__main__":
   assert strand_sort([4, 3, 5, 1, 2]) == [1, 2, 3, 4, 5]
   assert strand_sort([4, 3, 5, 1, 2], reverse=True) == [5, 4, 3, 2, 1]
