





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 							SCREAMING_SNAKE_CASE__						(    __a   , __a   ):
	snake_case_		:			str											=	[]
	for part_id in partition_order:
		snake_case_		:			List[str]											=	df.where(f"""SPARK_PARTITION_ID() = {part_id}"""   ).collect()
		for row_idx, row in enumerate(__a   ):
			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 							SCREAMING_SNAKE_CASE__						(    ):
	snake_case_		:			Tuple											=	pyspark.sql.SparkSession.builder.master('local[*]'   ).appName('pyspark'   ).getOrCreate()
	snake_case_		:			List[Any]											=	spark.range(1_00   ).repartition(1   )
	snake_case_		:			List[str]											=	Spark(__a   )
	# 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 							SCREAMING_SNAKE_CASE__						(    ):
	snake_case_		:			Union[str, Any]											=	pyspark.sql.SparkSession.builder.master('local[*]'   ).appName('pyspark'   ).getOrCreate()
	snake_case_		:			Dict											=	spark.range(10   ).repartition(2   )
	snake_case_		:			List[str]											=	[1, 0]
	snake_case_		:			str											=	_generate_iterable_examples(__a   , __a   )  # Reverse the partitions.
	snake_case_		:			Optional[int]											=	_get_expected_row_ids_and_row_dicts_for_partition_order(__a   , __a   )

	for i, (row_id, row_dict) in enumerate(generate_fn()   ):
		snake_case_						,snake_case_		:			str											=	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 							SCREAMING_SNAKE_CASE__						(    ):
	snake_case_		:			str											=	pyspark.sql.SparkSession.builder.master('local[*]'   ).appName('pyspark'   ).getOrCreate()
	snake_case_		:			Tuple											=	spark.range(10   ).repartition(1   )
	snake_case_		:			Optional[int]											=	SparkExamplesIterable(__a   )
	assert it.n_shards == 1
	for i, (row_id, row_dict) in enumerate(__a   ):
		assert row_id == f"""0_{i}"""
		assert row_dict == {"id": i}


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

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

	# Partitions 0 and 2
	snake_case_		:			List[str]											=	SparkExamplesIterable(__a   ).shard_data_sources(worker_id=0   , num_workers=2   )
	assert shard_it_a.n_shards == 2
	snake_case_		:			int											=	_get_expected_row_ids_and_row_dicts_for_partition_order(__a   , [0, 2]   )
	for i, (row_id, row_dict) in enumerate(__a   ):
		snake_case_						,snake_case_		:			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
	snake_case_		:			Union[str, Any]											=	SparkExamplesIterable(__a   ).shard_data_sources(worker_id=1   , num_workers=2   )
	assert shard_it_a.n_shards == 2
	snake_case_		:			List[str]											=	_get_expected_row_ids_and_row_dicts_for_partition_order(__a   , [1, 3]   )
	for i, (row_id, row_dict) in enumerate(__a   ):
		snake_case_						,snake_case_		:			Dict											=	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 							SCREAMING_SNAKE_CASE__						(    ):
	snake_case_		:			Union[str, Any]											=	pyspark.sql.SparkSession.builder.master('local[*]'   ).appName('pyspark'   ).getOrCreate()
	snake_case_		:			Dict											=	spark.range(1_00   ).repartition(1   )
	snake_case_		:			List[str]											=	Spark(__a   )
	# 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_00








from typing import TYPE_CHECKING

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


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

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

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

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_SCREAMING_SNAKE_CASE      	=		[
						    """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

						_SCREAMING_SNAKE_CASE      	=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

