'''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  _SCREAMING_SNAKE_CASE		(      __snake_case							:			Dict   ,		__snake_case							:			List[str]				):
     _A    		=  []
     for part_id in partition_order:
          _A    		=  df.where(F'SPARK_PARTITION_ID() = {part_id}'				).collect()
          for row_idx, row in enumerate(__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  _SCREAMING_SNAKE_CASE		(      ):
     _A    		=  pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
     _A    		=  spark.range(1_0_0				).repartition(1				)
     _A    		=  Spark(__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=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  _SCREAMING_SNAKE_CASE		(      ):
     _A    		=  pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
     _A    		=  spark.range(1_0				).repartition(2				)
     _A    		=  [1, 0]
     _A    		=  _generate_iterable_examples(__snake_case   ,		__snake_case				)  # Reverse the partitions.
     _A    		=  _get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case   ,		__snake_case				)

     for i, (row_id, row_dict) in enumerate(generate_fn()				):
          _A		,       _A    		=  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		(      ):
     _A    		=  pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
     _A    		=  spark.range(1_0				).repartition(1				)
     _A    		=  SparkExamplesIterable(__snake_case				)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(__snake_case				):
          assert row_id == F'0_{i}'
          assert row_dict == {"id": i}





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

          _A    		=  SparkExamplesIterable(__snake_case				).shuffle_data_sources(__snake_case				)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(__snake_case				):
               _A		,       _A    		=  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		(      ):
     _A    		=  pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
     _A    		=  spark.range(2_0				).repartition(4				)

     # Partitions 0 and 2
     _A    		=  SparkExamplesIterable(__snake_case				).shard_data_sources(worker_id=0   ,		num_workers=2				)
     assert shard_it_a.n_shards == 2
     _A    		=  _get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case   ,		[0, 2]				)
     for i, (row_id, row_dict) in enumerate(__snake_case				):
          _A		,       _A    		=  expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     _A    		=  SparkExamplesIterable(__snake_case				).shard_data_sources(worker_id=1   ,		num_workers=2				)
     assert shard_it_a.n_shards == 2
     _A    		=  _get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case   ,		[1, 3]				)
     for i, (row_id, row_dict) in enumerate(__snake_case				):
          _A		,       _A    		=  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		(      ):
     _A    		=  pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
     _A    		=  spark.range(1_0_0				).repartition(1				)
     _A    		=  Spark(__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() == 1_0_0




'''simple docstring'''






from typing import TYPE_CHECKING

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


_UpperCAmelCase    :     str         =       {
    '''configuration_vision_encoder_decoder''': ['''VisionEncoderDecoderConfig''', '''VisionEncoderDecoderOnnxConfig''']
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase    :     Tuple         =       ['''VisionEncoderDecoderModel''']

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase    :     str         =       ['''TFVisionEncoderDecoderModel''']

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase    :     Any         =       ['''FlaxVisionEncoderDecoderModel''']

if TYPE_CHECKING:
   from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig, VisionEncoderDecoderOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vision_encoder_decoder import VisionEncoderDecoderModel

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_vision_encoder_decoder import TFVisionEncoderDecoderModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_vision_encoder_decoder import FlaxVisionEncoderDecoderModel

else:
   import sys

   _UpperCAmelCase    :     List[str]         =       _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



