import logging
import os

from .state import PartialState




class  __magic_name__  (		logging.LoggerAdapter    ):





  """simple docstring"""




  @staticmethod
  def   lowerCAmelCase						(	_lowercase					:  Dict							):



   """simple docstring"""



   _UpperCamelCase:     List[str]							= PartialState()
   return not main_process_only or (main_process_only and state.is_main_process)



  def   lowerCAmelCase						(	self					:  List[Any]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  str		,  *_lowercase					:  int		,  **_lowercase					:  Union[str, Any]							):



   """simple docstring"""



   if PartialState._shared_state == {}:
    raise RuntimeError(
        '''You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility.'''							)
   _UpperCamelCase:     List[Any]							= kwargs.pop('''main_process_only'''		,  _lowercase							)
   _UpperCamelCase:     int							= kwargs.pop('''in_order'''		,  _lowercase							)

   if self.isEnabledFor(_lowercase							):
    if self._should_log(_lowercase							):
     _UpperCamelCase:     List[str]							= self.process(_lowercase		,  _lowercase							)
     self.logger.log(_lowercase		,  _lowercase		,  *_lowercase		,  **_lowercase							)

    elif in_order:
     _UpperCamelCase:     Union[str, Any]							= PartialState()
     for i in range(state.num_processes							):
      if i == state.process_index:
       _UpperCamelCase:     Any							= self.process(_lowercase		,  _lowercase							)
       self.logger.log(_lowercase		,  _lowercase		,  *_lowercase		,  **_lowercase							)
      state.wait_for_everyone()




def       lowerCAmelCase_     (		lowercase:							str  ,					lowercase:							str = None		)    ->				int:







 '''simple docstring'''
 if log_level is None:
  _UpperCamelCase:     Optional[int]							= os.environ.get('''ACCELERATE_LOG_LEVEL'''  ,					lowercase		)
 _UpperCamelCase:     Optional[Any]							= logging.getLogger(lowercase		)
 if log_level is not None:
  logger.setLevel(log_level.upper()		)
  logger.root.setLevel(log_level.upper()		)
 return MultiProcessAdapter(lowercase  ,					{}		)
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       lowerCAmelCase_     (		lowercase:							Optional[Any]  ,					lowercase:							Tuple		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     Union[str, Any]							= []
 for part_id in partition_order:
  _UpperCamelCase:     int							= df.where(F"""SPARK_PARTITION_ID() = {part_id}"""		).collect()
  for row_idx, row in enumerate(lowercase		):
   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       lowerCAmelCase_     (		)    ->				List[str]:







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







 '''simple docstring'''
 _UpperCamelCase:     List[Any]							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Optional[Any]							= spark.range(10		).repartition(2		)
 _UpperCamelCase:     int							= [1, 0]
 _UpperCamelCase:     Any							= _generate_iterable_examples(lowercase  ,					lowercase		)  # Reverse the partitions.
 _UpperCamelCase:     List[str]							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					lowercase		)

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







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




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				Union[str, Any]:







 '''simple docstring'''
 _UpperCamelCase:     Any							= 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:     Union[str, Any]							= lambda lowercase		: x.reverse()
  _UpperCamelCase:     List[str]							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					[2, 1, 0]		)

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







 '''simple docstring'''
 _UpperCamelCase:     str							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Tuple							= spark.range(20		).repartition(4		)

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







 '''simple docstring'''
 _UpperCamelCase:     Any							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Union[str, Any]							= spark.range(100		).repartition(1		)
 _UpperCamelCase:     Optional[Any]							= Spark(lowercase		)
 # 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