



'''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  :						Dict   ,      _lowerCAmelCase  :						Union[str, Any]  )      ->   Any:
							UpperCAmelCase :       str 							=		[]
							for part_id in partition_order:
														UpperCAmelCase :       List[str] 							=		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_     (					)      ->   List[Any]:
							UpperCAmelCase :       Optional[int] 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       str 							=		spark.range(100  ).repartition(1  )
							UpperCAmelCase :       Union[str, Any] 							=		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:
							UpperCAmelCase :       List[str] 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       str 							=		spark.range(10  ).repartition(2  )
							UpperCAmelCase :       Optional[int] 							=		[1, 0]
							UpperCAmelCase :       Tuple 							=		_generate_iterable_examples(_lowerCAmelCase   ,      _lowerCAmelCase  )  # Reverse the partitions.
							UpperCAmelCase :       Any 							=		_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_     (					)      ->   Union[str, Any]:
							UpperCAmelCase :       Dict 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       Optional[Any] 							=		spark.range(10  ).repartition(1  )
							UpperCAmelCase :       List[Any] 							=		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_     (					)      ->   Optional[Any]:
							UpperCAmelCase :       List[str] 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       str 							=		spark.range(30  ).repartition(3  )
							# Mock the generator so that shuffle reverses the partition indices.
							with patch('''numpy.random.Generator'''  ) as generator_mock:
														UpperCAmelCase :       List[str] 							=		lambda _lowerCAmelCase  : x.reverse()
														UpperCAmelCase :       Optional[int] 							=		_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCAmelCase   ,      [2, 1, 0]  )

														UpperCAmelCase :       str 							=		SparkExamplesIterable(_lowerCAmelCase  ).shuffle_data_sources(_lowerCAmelCase  )
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(_lowerCAmelCase  ):
																					UpperCAmelCase							,		UpperCAmelCase :       Dict 							=		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_     (					)      ->   Tuple:
							UpperCAmelCase :       Tuple 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       Any 							=		spark.range(20  ).repartition(4  )

							# Partitions 0 and 2
							UpperCAmelCase :       List[str] 							=		SparkExamplesIterable(_lowerCAmelCase  ).shard_data_sources(worker_id=0   ,      num_workers=2  )
							assert shard_it_a.n_shards == 2
							UpperCAmelCase :       Optional[Any] 							=		_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[int] 							=		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 :       str 							=		SparkExamplesIterable(_lowerCAmelCase  ).shard_data_sources(worker_id=1   ,      num_workers=2  )
							assert shard_it_a.n_shards == 2
							UpperCAmelCase :       List[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_     (					)      ->   List[Any]:
							UpperCAmelCase :       Optional[Any] 							=		pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
							UpperCAmelCase :       List[str] 							=		spark.range(100  ).repartition(1  )
							UpperCAmelCase :       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









import doctest
import logging
import os
import unittest
from pathlib import Path
from typing import List, Union

import transformers
from transformers.testing_utils import require_tf, require_torch, slow


lowerCamelCase											=       logging.getLogger()




@unittest.skip('Temporarily disable the doc tests.'     )
@require_torch
@require_tf
@slow
class A  (   unittest.TestCase     ):






      def        lowerCamelCase       (					self						:     Dict			,   lowercase_						:     Path			,   lowercase_						:     Union[str, None] = None			,   lowercase_						:     Union[List[str], None] = None			,   lowercase_						:     Union[str, List[str], None] = None			,   lowercase_						:     bool = True			,   )					->						Optional[int]:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=[file for file in os.listdir(lowercase_ ) if os.path.isfile(os.path.join(lowercase_			,   lowercase_ ) )]

             if identifier is not None:
                    _lowerCamelCase       :		Any		=[file for file in files if identifier in file]

             if n_identifier is not None:
                    if isinstance(lowercase_			,   lowercase_ ):
                           for n_ in n_identifier:
                                  _lowerCamelCase       :		Any		=[file for file in files if n_ not in file]
                    else:
                           _lowerCamelCase       :		List[Any]		=[file for file in files if n_identifier not in file]

             _lowerCamelCase       :		List[Any]		=ignore_files or []
             ignore_files.append('__init__.py' )
             _lowerCamelCase       :		List[Any]		=[file for file in files if file not in ignore_files]

             for file in files:
                    # Open all files
                    print('Testing'			,   lowercase_ )

                    if only_modules:
                           _lowerCamelCase       :		Optional[int]		=file.split('.' )[0]
                           try:
                                  _lowerCamelCase       :		Any		=getattr(lowercase_			,   lowercase_ )
                                  _lowerCamelCase       :		Optional[int]		=doctest.DocTestSuite(lowercase_ )
                                  _lowerCamelCase       :		List[Any]		=unittest.TextTestRunner().run(lowercase_ )
                                  self.assertIs(len(result.failures )			,   0 )
                           except AttributeError:
                                  logger.info(F'''{module_identifier} is not a module.''' )
                    else:
                           _lowerCamelCase       :		int		=doctest.testfile(str('..' / directory / file )			,   optionflags=doctest.ELLIPSIS )
                           self.assertIs(result.failed			,   0 )






      def        lowerCamelCase       (					self						:     Any )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		int		=Path('src/transformers' )
             _lowerCamelCase       :		Optional[Any]		='modeling'
             _lowerCamelCase       :		List[str]		=[
                 'modeling_ctrl.py',
                 'modeling_tf_ctrl.py',
             ]
             self.analyze_directory(lowercase_			,   identifier=lowercase_			,   ignore_files=lowercase_ )






      def        lowerCamelCase       (					self						:     List[str] )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=Path('src/transformers' )
             _lowerCamelCase       :		Union[str, Any]		='tokenization'
             self.analyze_directory(lowercase_			,   identifier=lowercase_ )






      def        lowerCamelCase       (					self						:     Optional[int] )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		str		=Path('src/transformers' )
             _lowerCamelCase       :		Any		='configuration'
             self.analyze_directory(lowercase_			,   identifier=lowercase_ )






      def        lowerCamelCase       (					self						:     List[str] )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		Union[str, Any]		=Path('src/transformers' )
             _lowerCamelCase       :		int		=['configuration', 'modeling', 'tokenization']
             self.analyze_directory(lowercase_			,   n_identifier=lowercase_ )






      def        lowerCamelCase       (					self						:     Optional[Any] )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		str		=Path('docs/source' )
             _lowerCamelCase       :		Dict		=['favicon.ico']
             self.analyze_directory(lowercase_			,   ignore_files=lowercase_			,   only_modules=lowercase_ )


