





"""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  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )   ->Optional[int]:





  """simple docstring"""


  lowerCAmelCase__     :int											=				[]
  for part_id in partition_order:
    lowerCAmelCase__     :Dict											=				df.where(F"SPARK_PARTITION_ID() = {part_id}"     ).collect()
    for row_idx, row in enumerate(_SCREAMING_SNAKE_CASE     ):
      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  ()   ->int:





  """simple docstring"""


  lowerCAmelCase__     :Tuple											=				pyspark.sql.SparkSession.builder.master('local[*]'     ).appName('pyspark'     ).getOrCreate()
  lowerCAmelCase__     :Union[str, Any]											=				spark.range(100     ).repartition(1     )
  lowerCAmelCase__     :int											=				Spark(_SCREAMING_SNAKE_CASE     )
  # 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 					__A  ()   ->Dict:





  """simple docstring"""


  lowerCAmelCase__     :Tuple											=				pyspark.sql.SparkSession.builder.master('local[*]'     ).appName('pyspark'     ).getOrCreate()
  lowerCAmelCase__     :Dict											=				spark.range(10     ).repartition(2     )
  lowerCAmelCase__     :Union[str, Any]											=				[1, 0]
  lowerCAmelCase__     :int											=				_generate_iterable_examples(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )  # Reverse the partitions.
  lowerCAmelCase__     :Dict											=				_get_expected_row_ids_and_row_dicts_for_partition_order(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )

  for i, (row_id, row_dict) in enumerate(generate_fn()     ):
    lowerCAmelCase__				,							lowerCAmelCase__     :Optional[int]											=				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"""


  lowerCAmelCase__     :Dict											=				pyspark.sql.SparkSession.builder.master('local[*]'     ).appName('pyspark'     ).getOrCreate()
  lowerCAmelCase__     :int											=				spark.range(10     ).repartition(1     )
  lowerCAmelCase__     :Any											=				SparkExamplesIterable(_SCREAMING_SNAKE_CASE     )
  assert it.n_shards == 1
  for i, (row_id, row_dict) in enumerate(_SCREAMING_SNAKE_CASE     ):
    assert row_id == F"0_{i}"
    assert row_dict == {"id": i}
@require_not_windows
@require_dill_gt_0_3_2
def 					__A  ()   ->Optional[Any]:





  """simple docstring"""


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

    lowerCAmelCase__     :Any											=				SparkExamplesIterable(_SCREAMING_SNAKE_CASE     ).shuffle_data_sources(_SCREAMING_SNAKE_CASE     )
    assert shuffled_it.n_shards == 3
    for i, (row_id, row_dict) in enumerate(_SCREAMING_SNAKE_CASE     ):
      lowerCAmelCase__				,							lowerCAmelCase__     :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  ()   ->Optional[int]:





  """simple docstring"""


  lowerCAmelCase__     :Dict											=				pyspark.sql.SparkSession.builder.master('local[*]'     ).appName('pyspark'     ).getOrCreate()
  lowerCAmelCase__     :str											=				spark.range(20     ).repartition(4     )

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





  """simple docstring"""


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


_snake_case : str  				=						{
    'configuration_layoutlmv3': [
        'LAYOUTLMV3_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'LayoutLMv3Config',
        'LayoutLMv3OnnxConfig',
    ],
    'processing_layoutlmv3': ['LayoutLMv3Processor'],
    'tokenization_layoutlmv3': ['LayoutLMv3Tokenizer'],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_snake_case : List[str]  				=						['LayoutLMv3TokenizerFast']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_snake_case : Optional[int]  				=						[
				    'LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'LayoutLMv3ForQuestionAnswering',
				    'LayoutLMv3ForSequenceClassification',
				    'LayoutLMv3ForTokenClassification',
				    'LayoutLMv3Model',
				    'LayoutLMv3PreTrainedModel',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_snake_case : Tuple  				=						[
				    'TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'TFLayoutLMv3ForQuestionAnswering',
				    'TFLayoutLMv3ForSequenceClassification',
				    'TFLayoutLMv3ForTokenClassification',
				    'TFLayoutLMv3Model',
				    'TFLayoutLMv3PreTrainedModel',
				]

try:
				if not is_vision_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_snake_case : List[Any]  				=						['LayoutLMv3FeatureExtractor']
				_snake_case : Tuple  				=						['LayoutLMv3ImageProcessor']


if TYPE_CHECKING:
				from .configuration_layoutlmva import (
				    LAYOUTLMV3_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    LayoutLMvaConfig,
				    LayoutLMvaOnnxConfig,
				)
				from .processing_layoutlmva import LayoutLMvaProcessor
				from .tokenization_layoutlmva import LayoutLMvaTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_layoutlmva_fast import LayoutLMvaTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_layoutlmva import (
								    LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
								    LayoutLMvaForQuestionAnswering,
								    LayoutLMvaForSequenceClassification,
								    LayoutLMvaForTokenClassification,
								    LayoutLMvaModel,
								    LayoutLMvaPreTrainedModel,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_layoutlmva import (
								    TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFLayoutLMvaForQuestionAnswering,
								    TFLayoutLMvaForSequenceClassification,
								    TFLayoutLMvaForTokenClassification,
								    TFLayoutLMvaModel,
								    TFLayoutLMvaPreTrainedModel,
								)

				try:
								if not is_vision_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .feature_extraction_layoutlmva import LayoutLMvaFeatureExtractor
								from .image_processing_layoutlmva import LayoutLMvaImageProcessor

else:
				import sys

				_snake_case : Any  				=						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

