'''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 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Dict:




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




						'''simple docstring'''
						_UpperCamelCase			:			str   			=						pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
						_UpperCamelCase			:			int   			=						spark.range(1_0_0					).repartition(1					)
						_UpperCamelCase			:			str   			=						Spark(UpperCAmelCase					)
						# 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 						__A							(     )     ->			Dict:




						'''simple docstring'''
						_UpperCamelCase			:			List[Any]   			=						pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
						_UpperCamelCase			:			Dict   			=						spark.range(1_0					).repartition(2					)
						_UpperCamelCase			:			Union[str, Any]   			=						[1, 0]
						_UpperCamelCase			:			Dict   			=						_generate_iterable_examples(UpperCAmelCase     ,UpperCAmelCase					)  # Reverse the partitions.
						_UpperCamelCase			:			Tuple   			=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase     ,UpperCAmelCase					)

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




						'''simple docstring'''
						_UpperCamelCase			:			List[Any]   			=						pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
						_UpperCamelCase			:			Union[str, Any]   			=						spark.range(1_0					).repartition(1					)
						_UpperCamelCase			:			Any   			=						SparkExamplesIterable(UpperCAmelCase					)
						assert it.n_shards == 1
						for i, (row_id, row_dict) in enumerate(UpperCAmelCase					):
												assert row_id == f'''0_{i}'''
												assert row_dict == {"id": i}
@require_not_windows
@require_dill_gt_0_3_2
def 						__A							(     )     ->			str:




						'''simple docstring'''
						_UpperCamelCase			:			List[Any]   			=						pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
						_UpperCamelCase			:			Tuple   			=						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 UpperCAmelCase					: x.reverse()
												_UpperCamelCase			:			List[str]   			=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase     ,[2, 1, 0]					)

												_UpperCamelCase			:			Any   			=						SparkExamplesIterable(UpperCAmelCase					).shuffle_data_sources(UpperCAmelCase					)
												assert shuffled_it.n_shards == 3
												for i, (row_id, row_dict) in enumerate(UpperCAmelCase					):
																		_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 						__A							(     )     ->			str:




						'''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			:			str   			=						SparkExamplesIterable(UpperCAmelCase					).shard_data_sources(worker_id=0     ,num_workers=2					)
						assert shard_it_a.n_shards == 2
						_UpperCamelCase			:			str   			=						_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCAmelCase     ,[0, 2]					)
						for i, (row_id, row_dict) in enumerate(UpperCAmelCase					):
												_UpperCamelCase			,       _UpperCamelCase			:			List[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			:			Optional[Any]   			=						SparkExamplesIterable(UpperCAmelCase					).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(UpperCAmelCase     ,[1, 3]					)
						for i, (row_id, row_dict) in enumerate(UpperCAmelCase					):
												_UpperCamelCase			,       _UpperCamelCase			:			List[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 						__A							(     )     ->			List[Any]:




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



'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCAmelCase_	:		Union[str, Any]			=     {"""configuration_opt""": ["""OPT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OPTConfig"""]}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		Tuple			=     [
			    """OPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
			    """OPTForCausalLM""",
			    """OPTModel""",
			    """OPTPreTrainedModel""",
			    """OPTForSequenceClassification""",
			    """OPTForQuestionAnswering""",
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		str			=     ["""TFOPTForCausalLM""", """TFOPTModel""", """TFOPTPreTrainedModel"""]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		List[Any]			=     [
			    """FlaxOPTForCausalLM""",
			    """FlaxOPTModel""",
			    """FlaxOPTPreTrainedModel""",
			]


if TYPE_CHECKING:
			from .configuration_opt import OPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OPTConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_opt import (
						    OPT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    OPTForCausalLM,
						    OPTForQuestionAnswering,
						    OPTForSequenceClassification,
						    OPTModel,
						    OPTPreTrainedModel,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_opt import TFOPTForCausalLM, TFOPTModel, TFOPTPreTrainedModel

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_opt import FlaxOPTForCausalLM, FlaxOPTModel, FlaxOPTPreTrainedModel

else:
			import sys

			lowerCAmelCase_	:		Any			=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


