



'''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










'''simple docstring'''






from typing import TYPE_CHECKING

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


UpperCamelCase__:   str								=			{
    "configuration_lxmert": ["LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "LxmertConfig"],
    "tokenization_lxmert": ["LxmertTokenizer"],
}

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__:   int								=			["LxmertTokenizerFast"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__:   Union[str, Any]								=			[
						    "LxmertEncoder",
						    "LxmertForPreTraining",
						    "LxmertForQuestionAnswering",
						    "LxmertModel",
						    "LxmertPreTrainedModel",
						    "LxmertVisualFeatureEncoder",
						    "LxmertXLayer",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__:   int								=			[
						    "TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "TFLxmertForPreTraining",
						    "TFLxmertMainLayer",
						    "TFLxmertModel",
						    "TFLxmertPreTrainedModel",
						    "TFLxmertVisualFeatureEncoder",
						]


if TYPE_CHECKING:
						from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
						from .tokenization_lxmert import LxmertTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_lxmert_fast import LxmertTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_lxmert import (
												    LxmertEncoder,
												    LxmertForPreTraining,
												    LxmertForQuestionAnswering,
												    LxmertModel,
												    LxmertPreTrainedModel,
												    LxmertVisualFeatureEncoder,
												    LxmertXLayer,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_lxmert import (
												    TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFLxmertForPreTraining,
												    TFLxmertMainLayer,
												    TFLxmertModel,
												    TFLxmertPreTrainedModel,
												    TFLxmertVisualFeatureEncoder,
												)

else:
						import sys

						UpperCamelCase__:   Tuple								=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





