






'''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    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[Any],					SCREAMING_SNAKE_CASE__:	Any    )  ->			Optional[int]:

   """simple docstring"""
   __a				     =    []
   for part_id in partition_order:
      __a				     =    df.where(f"""SPARK_PARTITION_ID() = {part_id}"""    ).collect()
      for row_idx, row in enumerate(SCREAMING_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    __UpperCAmelCase	(				)  ->			Optional[int]:

   """simple docstring"""
   __a				     =    pyspark.sql.SparkSession.builder.master('local[*]'    ).appName('pyspark'    ).getOrCreate()
   __a				     =    spark.range(100    ).repartition(1    )
   __a				     =    Spark(SCREAMING_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=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	(				)  ->			Optional[int]:

   """simple docstring"""
   __a				     =    pyspark.sql.SparkSession.builder.master('local[*]'    ).appName('pyspark'    ).getOrCreate()
   __a				     =    spark.range(10    ).repartition(2    )
   __a				     =    [1, 0]
   __a				     =    _generate_iterable_examples(SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__    )  # Reverse the partitions.
   __a				     =    _get_expected_row_ids_and_row_dicts_for_partition_order(SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__    )

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

   """simple docstring"""
   __a				     =    pyspark.sql.SparkSession.builder.master('local[*]'    ).appName('pyspark'    ).getOrCreate()
   __a				     =    spark.range(10    ).repartition(1    )
   __a				     =    SparkExamplesIterable(SCREAMING_SNAKE_CASE__    )
   assert it.n_shards == 1
   for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__    ):
      assert row_id == f"""0_{i}"""
      assert row_dict == {"id": i}



@require_not_windows
@require_dill_gt_0_3_2
def    __UpperCAmelCase	(				)  ->			Optional[int]:

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

      __a				     =    SparkExamplesIterable(SCREAMING_SNAKE_CASE__    ).shuffle_data_sources(SCREAMING_SNAKE_CASE__    )
      assert shuffled_it.n_shards == 3
      for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__    ):
         __a     ,			__a				     =    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	(				)  ->			Tuple:

   """simple docstring"""
   __a				     =    pyspark.sql.SparkSession.builder.master('local[*]'    ).appName('pyspark'    ).getOrCreate()
   __a				     =    spark.range(20    ).repartition(4    )

   # Partitions 0 and 2
   __a				     =    SparkExamplesIterable(SCREAMING_SNAKE_CASE__    ).shard_data_sources(worker_id=0,					num_workers=2    )
   assert shard_it_a.n_shards == 2
   __a				     =    _get_expected_row_ids_and_row_dicts_for_partition_order(SCREAMING_SNAKE_CASE__,					[0, 2]    )
   for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__    ):
      __a     ,			__a				     =    expected_row_ids_and_row_dicts_a[i]
      assert row_id == expected_row_id
      assert row_dict == expected_row_dict

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

   """simple docstring"""
   __a				     =    pyspark.sql.SparkSession.builder.master('local[*]'    ).appName('pyspark'    ).getOrCreate()
   __a				     =    spark.range(100    ).repartition(1    )
   __a				     =    Spark(SCREAMING_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() == 100







'''simple docstring'''



import warnings

from .generation import TFGenerationMixin
class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):
      # warning at import time
      warnings.warn(
          "Importing `TFGenerationMixin` from `src/transformers/generation_tf_utils.py` is deprecated and will "
          "be removed in Transformers v5. Import as `from transformers import TFGenerationMixin` instead."      ,      _lowerCAmelCase      ,      )