



"""simple docstring"""

import io
import math
from typing import Dict, Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import convert_to_rgb, normalize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    get_image_size,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_vision_available, logging
from ...utils.import_utils import requires_backends


if is_vision_available():
     import textwrap

     from PIL import Image, ImageDraw, ImageFont

if is_torch_available():
     import torch

     from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
     lowercase__    :  Dict							    =	False


lowercase__    :  Any							    =	logging.get_logger(__name__)
lowercase__    :  List[Any]							    =	'''ybelkada/fonts'''




def __lowercase			(							):
    if is_torch_available() and not is_torch_greater_or_equal_than_1_11:
        raise ImportError(
            f"You are using torch=={torch.__version__}, but torch>=1.11.0 is required to use "
            '''Pix2StructImageProcessor. Please upgrade torch.'''    )




def __lowercase			(							_a		,			_a		,			_a    ):
    requires_backends(_a		,			['''torch''']    )
    _check_torch_version()

    snake_case_ :     Optional[int]				=	image_tensor.unsqueeze(0    )
    snake_case_ :     List[str]				=	torch.nn.functional.unfold(_a		,			(patch_height, patch_width)		,			stride=(patch_height, patch_width)    )
    snake_case_ :     Union[str, Any]				=	patches.reshape(image_tensor.size(0    )		,			image_tensor.size(1    )		,			_a		,			_a		,			-1    )
    snake_case_ :     List[Any]				=	patches.permute(0		,			4		,			2		,			3		,			1    ).reshape(
        image_tensor.size(2    ) // patch_height		,			image_tensor.size(3    ) // patch_width		,			image_tensor.size(1    ) * patch_height * patch_width		,			)
    return patches.unsqueeze(0    )




def __lowercase			(							_a		,			_a = 36		,			_a = "black"		,			_a = "white"		,			_a = 5		,			_a = 5		,			_a = 5		,			_a = 5		,			_a = None		,			_a = None		,			):
    requires_backends(_a		,			'''vision'''    )
    # Add new lines so that each line is no more than 80 characters.

    snake_case_ :     Dict				=	textwrap.TextWrapper(width=80    )
    snake_case_ :     List[str]				=	wrapper.wrap(text=_a    )
    snake_case_ :     Tuple				=	'\n'.join(_a    )

    if font_bytes is not None and font_path is None:
        snake_case_ :     Any				=	io.BytesIO(_a    )
    elif font_path is not None:
        snake_case_ :     str				=	font_path
    else:
        snake_case_ :     Any				=	hf_hub_download(_a		,			'''Arial.TTF'''    )
    snake_case_ :     Optional[int]				=	ImageFont.truetype(_a		,			encoding='''UTF-8'''		,			size=_a    )

    # Use a temporary canvas to determine the width and height in pixels when
    # rendering the text.
    snake_case_ :     int				=	ImageDraw.Draw(Image.new('''RGB'''		,			(1, 1)		,			_a    )    )
    snake_case_ :     Optional[Any]				=	temp_draw.textbbox((0, 0)		,			_a		,			_a    )

    # Create the actual image with a bit of padding around the text.
    snake_case_ :     Optional[Any]				=	text_width + left_padding + right_padding
    snake_case_ :     Tuple				=	text_height + top_padding + bottom_padding
    snake_case_ :     Union[str, Any]				=	Image.new('''RGB'''		,			(image_width, image_height)		,			_a    )
    snake_case_ :     List[Any]				=	ImageDraw.Draw(_a    )
    draw.text(xy=(left_padding, top_padding)		,			text=_a		,			fill=_a		,			font=_a    )
    return image




def __lowercase			(							_a		,			_a		,			**_a    ):
    requires_backends(_a		,			'''vision'''    )

    # Convert to PIL image if necessary
    snake_case_ :     Tuple				=	to_pil_image(_a    )

    snake_case_ :     Union[str, Any]				=	render_text(_a		,			**_a    )
    snake_case_ :     List[Any]				=	max(header_image.width		,			image.width    )

    snake_case_ :     Any				=	int(image.height * (new_width / image.width)    )
    snake_case_ :     Union[str, Any]				=	int(header_image.height * (new_width / header_image.width)    )

    snake_case_ :     List[Any]				=	Image.new('''RGB'''		,			(new_width, new_height + new_header_height)		,			'''white'''    )
    new_image.paste(header_image.resize((new_width, new_header_height)    )		,			(0, 0)    )
    new_image.paste(image.resize((new_width, new_height)    )		,			(0, new_header_height)    )

    # Convert back to the original framework if necessary
    snake_case_ :     Optional[Any]				=	to_numpy_array(_a    )

    if infer_channel_dimension_format(_a    ) == ChannelDimension.LAST:
        snake_case_ :     int				=	to_channel_dimension_format(_a		,			ChannelDimension.LAST    )

    return new_image






class        _UpperCAmelCase       (			snake_case_):
 _lowerCAmelCase   :			str							=		["""flattened_patches"""]

 def __init__(     self		:       int  ,     lowercase_		:       bool = True  ,     lowercase_		:       bool = True  ,     lowercase_		:       Dict[str, int] = None  ,     lowercase_		:       int = 2048  ,     lowercase_		:       bool = False  ,     **lowercase_		:       str  ,     ):
     super().__init__(**lowercase_				)
     snake_case_ :     Union[str, Any]				=	patch_size if patch_size is not None else {'height': 16, 'width': 16}
     snake_case_ :     Dict				=	do_normalize
     snake_case_ :     int				=	do_convert_rgb
     snake_case_ :     List[str]				=	max_patches
     snake_case_ :     List[str]				=	is_vqa

 def 						_snake_case					(     self		:       Union[str, Any]  ,     lowercase_		:       np.ndarray  ,     lowercase_		:       int  ,     lowercase_		:       dict  ,     **lowercase_		:       Dict				):
     requires_backends(self.extract_flattened_patches  ,     '''torch'''				)
     _check_torch_version()

     # convert to torch
     snake_case_ :     Optional[Any]				=	to_channel_dimension_format(lowercase_  ,     ChannelDimension.FIRST				)
     snake_case_ :     str				=	torch.from_numpy(lowercase_				)

     snake_case_ :     Optional[Any]				=	patch_size['height'], patch_size['width']
     snake_case_ :     List[str]				=	get_image_size(lowercase_				)

     # maximize scale s.t.
     snake_case_ :     Any				=	math.sqrt(max_patches * (patch_height / image_height) * (patch_width / image_width)				)
     snake_case_ :     Optional[int]				=	max(min(math.floor(scale * image_height / patch_height				)  ,     lowercase_				)  ,     1				)
     snake_case_ :     str				=	max(min(math.floor(scale * image_width / patch_width				)  ,     lowercase_				)  ,     1				)
     snake_case_ :     List[Any]				=	max(num_feasible_rows * patch_height  ,     1				)
     snake_case_ :     List[str]				=	max(num_feasible_cols * patch_width  ,     1				)

     snake_case_ :     Any				=	torch.nn.functional.interpolate(
         image.unsqueeze(0				)  ,     size=(resized_height, resized_width)  ,     mode='''bilinear'''  ,     align_corners=lowercase_  ,     antialias=lowercase_  ,     ).squeeze(0				)

     # [1, rows, columns, patch_height * patch_width * image_channels]
     snake_case_ :     Union[str, Any]				=	torch_extract_patches(lowercase_  ,     lowercase_  ,     lowercase_				)

     snake_case_ :     str				=	patches.shape
     snake_case_ :     List[Any]				=	patches_shape[1]
     snake_case_ :     Union[str, Any]				=	patches_shape[2]
     snake_case_ :     Optional[Any]				=	patches_shape[3]

     # [rows * columns, patch_height * patch_width * image_channels]
     snake_case_ :     Tuple				=	patches.reshape([rows * columns, depth]				)

     # [rows * columns, 1]
     snake_case_ :     Union[str, Any]				=	torch.arange(lowercase_				).reshape([rows, 1]				).repeat(1  ,     lowercase_				).reshape([rows * columns, 1]				)
     snake_case_ :     List[Any]				=	torch.arange(lowercase_				).reshape([1, columns]				).repeat(lowercase_  ,     1				).reshape([rows * columns, 1]				)

     # Offset by 1 so the ids do not contain zeros, which represent padding.
     row_ids += 1
     col_ids += 1

     # Prepare additional patch features.
     # [rows * columns, 1]
     snake_case_ :     Any				=	row_ids.to(torch.floataa				)
     snake_case_ :     Any				=	col_ids.to(torch.floataa				)

     # [rows * columns, 2 + patch_height * patch_width * image_channels]
     snake_case_ :     Dict				=	torch.cat([row_ids, col_ids, patches]  ,     -1				)

     # [max_patches, 2 + patch_height * patch_width * image_channels]
     snake_case_ :     Tuple				=	torch.nn.functional.pad(lowercase_  ,     [0, 0, 0, max_patches - (rows * columns)]				).float()

     snake_case_ :     List[Any]				=	to_numpy_array(lowercase_				)

     return result

 def 						_snake_case					(     self		:       Any  ,     lowercase_		:       np.ndarray  ,     lowercase_		:       Optional[Union[str, ChannelDimension]] = None  ,     **lowercase_		:       Dict				):





     if image.dtype == np.uinta:
         snake_case_ :     int				=	image.astype(np.floataa				)

     # take mean across the whole `image`
     snake_case_ :     Optional[Any]				=	np.mean(lowercase_				)
     snake_case_ :     List[str]				=	np.std(lowercase_				)
     snake_case_ :     Any				=	max(lowercase_  ,     1.0 / math.sqrt(np.prod(image.shape				)				)				)

     return normalize(lowercase_  ,     mean=lowercase_  ,     std=lowercase_  ,     **lowercase_				)






 def 						_snake_case					(     self		:       Optional[int]  ,     lowercase_		:       ImageInput  ,     lowercase_		:       Optional[str] = None  ,     lowercase_		:       bool = None  ,     lowercase_		:       Optional[bool] = None  ,     lowercase_		:       Optional[int] = None  ,     lowercase_		:       Optional[Dict[str, int]] = None  ,     lowercase_		:       Optional[Union[str, TensorType]] = None  ,     lowercase_		:       ChannelDimension = ChannelDimension.FIRST  ,     **lowercase_		:       str  ,     ):
     snake_case_ :     int				=	do_normalize if do_normalize is not None else self.do_normalize
     snake_case_ :     List[str]				=	do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
     snake_case_ :     Any				=	patch_size if patch_size is not None else self.patch_size
     snake_case_ :     List[Any]				=	max_patches if max_patches is not None else self.max_patches
     snake_case_ :     Optional[Any]				=	self.is_vqa

     if kwargs.get('''data_format'''  ,     lowercase_				) is not None:
         raise ValueError('''data_format is not an accepted input as the outputs are '''				)

     snake_case_ :     List[str]				=	make_list_of_images(lowercase_				)

     if not valid_images(lowercase_				):
         raise ValueError(
             '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
             '''torch.Tensor, tf.Tensor or jax.ndarray.'''				)

     # PIL RGBA images are converted to RGB
     if do_convert_rgb:
         snake_case_ :     List[Any]				=	[convert_to_rgb(lowercase_				) for image in images]

     # All transformations expect numpy arrays.
     snake_case_ :     Union[str, Any]				=	[to_numpy_array(lowercase_				) for image in images]

     if is_vqa:
         if header_text is None:
             raise ValueError('''A header text must be provided for VQA models.'''				)
         snake_case_ :     Optional[int]				=	kwargs.pop('''font_bytes'''  ,     lowercase_				)
         snake_case_ :     List[str]				=	kwargs.pop('''font_path'''  ,     lowercase_				)

         if isinstance(lowercase_  ,     lowercase_				):
             snake_case_ :     Optional[int]				=	[header_text] * len(lowercase_				)

         snake_case_ :     Union[str, Any]				=	[
             render_header(lowercase_  ,     header_text[i]  ,     font_bytes=lowercase_  ,     font_path=lowercase_				)
             for i, image in enumerate(lowercase_				)
         ]

     if do_normalize:
         snake_case_ :     Optional[int]				=	[self.normalize(image=lowercase_				) for image in images]

     # convert to torch tensor and permute
     snake_case_ :     List[str]				=	[
         self.extract_flattened_patches(image=lowercase_  ,     max_patches=lowercase_  ,     patch_size=lowercase_				)
         for image in images
     ]

     # create attention mask in numpy
     snake_case_ :     int				=	[(image.sum(axis=-1				) != 0).astype(np.floataa				) for image in images]

     snake_case_ :     Optional[Any]				=	BatchFeature(
         data={'''flattened_patches''': images, '''attention_mask''': attention_masks}  ,     tensor_type=lowercase_				)

     return encoded_outputs






import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


a_    					=       datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
   import pyspark





@dataclass
class 		_lowercase						(					datasets.BuilderConfig       ):
 lowercase				     =   None





def    __lowercase     (						lowerCamelCase					:     "pyspark.sql.DataFrame"  ,  lowerCamelCase					:     List[int]  ,  ):
 import pyspark

 def generate_fn():
  UpperCamelCase_					:						Dict =					df.select('*'  ,  pyspark.sql.functions.spark_partition_id().alias('part_id'	)	)
  for partition_id in partition_order:
   UpperCamelCase_					:						Tuple =					df_with_partition_id.select('*'	).where(F"part_id = {partition_id}"	).drop('part_id'	)
   UpperCamelCase_					:						Union[str, Any] =					partition_df.collect()
   UpperCamelCase_					:						Any =					0
   for row in rows:
    yield F"{partition_id}_{row_id}", row.asDict()
    row_id += 1

 return generate_fn





class 		_lowercase						(					_BaseExamplesIterable       ):

 def __init__( self	:							Optional[int]   ,      snake_case	:							"pyspark.sql.DataFrame"   ,      snake_case	:							Tuple=None   ,      )   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_					:						Dict =					df
  UpperCamelCase_					:						int =					partition_order or range(self.df.rdd.getNumPartitions()			)
  UpperCamelCase_					:						Optional[Any] =					_generate_iterable_examples(self.df   ,      self.partition_order			)

 def __iter__( self	:							Optional[int]			)   ->						Any:

  """simple docstring"""





  yield from self.generate_examples_fn()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any   ,      snake_case	:							np.random.Generator			)   ->						"SparkExamplesIterable":

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					list(range(self.df.rdd.getNumPartitions()			)			)
  generator.shuffle(snake_case			)
  return SparkExamplesIterable(self.df   ,      partition_order=snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int   ,      snake_case	:							int   ,      snake_case	:							int			)   ->						"SparkExamplesIterable":

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					self.split_shard_indices_by_worker(snake_case   ,      snake_case			)
  return SparkExamplesIterable(self.df   ,      partition_order=snake_case			)


 @property
 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						int:

  """simple docstring"""





  return len(self.partition_order			)





class 		_lowercase						(					datasets.DatasetBuilder       ):
 lowercase				     =   SparkConfig

 def __init__( self	:							List[Any]   ,      snake_case	:							"pyspark.sql.DataFrame"   ,      snake_case	:							str = None   ,      snake_case	:							str = None   ,      **snake_case	:							Optional[Any]   ,      )   ->						List[str]:

  """simple docstring"""





  import pyspark

  UpperCamelCase_					:						List[Any] =					pyspark.sql.SparkSession.builder.getOrCreate()
  UpperCamelCase_					:						str =					df
  UpperCamelCase_					:						Tuple =					working_dir

  super().__init__(
      cache_dir=snake_case   ,      config_name=str(self.df.semanticHash()			)   ,      **snake_case   ,      )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						Dict:

  """simple docstring"""





  def create_cache_and_write_probe(snake_case	:							str			):
   # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
   # already exist.
   os.makedirs(self._cache_dir   ,      exist_ok=snake_case			)
   UpperCamelCase_					:						Tuple =					os.path.join(self._cache_dir   ,      'fs_test' + uuid.uuida().hex			)
   # Opening the file in append mode will create a new file unless it already exists, in which case it will not
   # change the file contents.
   open(snake_case   ,      'a'			)
   return [probe_file]

  if self._spark.conf.get('spark.master'   ,      ''			).startswith('local'			):
   return

  # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
  # accessible to the driver.
  # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
  if self._cache_dir:
   UpperCamelCase_					:						Tuple =					(
       self._spark.sparkContext.parallelize(range(1			)   ,      1			).mapPartitions(snake_case			).collect()
   )
   if os.path.isfile(probe[0]			):
    return

  raise ValueError(
      'When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						Union[str, Any]:

  """simple docstring"""





  return datasets.DatasetInfo(features=self.config.features			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							datasets.download.download_manager.DownloadManager			)   ->						Optional[int]:

  """simple docstring"""





  return [datasets.SplitGenerator(name=datasets.Split.TRAIN			)]

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							Optional[int]			)   ->						List[Any]:

  """simple docstring"""





  import pyspark

  def get_arrow_batch_size(snake_case	:							Dict			):
   for batch in it:
    yield pa.RecordBatch.from_pydict({'batch_bytes': [batch.nbytes]}			)

  UpperCamelCase_					:						List[str] =					self.df.count()
  UpperCamelCase_					:						Union[str, Any] =					df_num_rows if df_num_rows <= 1_0_0 else 1_0_0
  # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
  UpperCamelCase_					:						str =					(
      self.df.limit(snake_case			)
      .repartition(1			)
      .mapInArrow(snake_case   ,      'batch_bytes: long'			)
      .agg(pyspark.sql.functions.sum('batch_bytes'			).alias('sample_bytes'			)			)
      .collect()[0]
      .sample_bytes
      / sample_num_rows
  )
  UpperCamelCase_					:						Optional[int] =					approx_bytes_per_row * df_num_rows
  if approx_total_size > max_shard_size:
   # Make sure there is at least one row per partition.
   UpperCamelCase_					:						Optional[Any] =					min(snake_case   ,      int(approx_total_size / max_shard_size			)			)
   UpperCamelCase_					:						int =					self.df.repartition(snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]   ,      snake_case	:							str   ,      snake_case	:							str   ,      snake_case	:							int   ,      )   ->						Iterable[Tuple[int, bool, Union[int, tuple]]]:

  """simple docstring"""





  import pyspark

  UpperCamelCase_					:						List[Any] =					ParquetWriter if file_format == 'parquet' else ArrowWriter
  UpperCamelCase_					:						List[str] =					os.path.join(self._working_dir   ,      os.path.basename(snake_case			)			) if self._working_dir else fpath
  UpperCamelCase_					:						Union[str, Any] =					file_format == 'parquet'

  # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
  # pickling the SparkContext.
  UpperCamelCase_					:						Union[str, Any] =					self.config.features
  UpperCamelCase_					:						Any =					self._writer_batch_size
  UpperCamelCase_					:						Dict =					self._fs.storage_options

  def write_arrow(snake_case	:							List[str]			):
   # Within the same SparkContext, no two task attempts will share the same attempt ID.
   UpperCamelCase_					:						Any =					pyspark.TaskContext().taskAttemptId()
   UpperCamelCase_					:						str =					next(snake_case   ,      snake_case			)
   if first_batch is None:
    # Some partitions might not receive any data.
    return pa.RecordBatch.from_arrays(
        [[task_id], [0], [0]]   ,      names=['task_id', 'num_examples', 'num_bytes']   ,      )
   UpperCamelCase_					:						Any =					0
   UpperCamelCase_					:						Optional[Any] =					writer_class(
       features=snake_case   ,      path=working_fpath.replace('SSSSS'   ,      f"{shard_id:05d}"			).replace('TTTTT'   ,      f"{task_id:05d}"			)   ,      writer_batch_size=snake_case   ,      storage_options=snake_case   ,      embed_local_files=snake_case   ,      )
   UpperCamelCase_					:						str =					pa.Table.from_batches([first_batch]			)
   writer.write_table(snake_case			)
   for batch in it:
    if max_shard_size is not None and writer._num_bytes >= max_shard_size:
     UpperCamelCase_, UpperCamelCase_					:						str =					writer.finalize()
     writer.close()
     yield pa.RecordBatch.from_arrays(
         [[task_id], [num_examples], [num_bytes]]   ,      names=['task_id', 'num_examples', 'num_bytes']   ,      )
     shard_id += 1
     UpperCamelCase_					:						Union[str, Any] =					writer_class(
         features=writer._features   ,      path=working_fpath.replace('SSSSS'   ,      f"{shard_id:05d}"			).replace('TTTTT'   ,      f"{task_id:05d}"			)   ,      writer_batch_size=snake_case   ,      storage_options=snake_case   ,      embed_local_files=snake_case   ,      )
    UpperCamelCase_					:						Optional[Any] =					pa.Table.from_batches([batch]			)
    writer.write_table(snake_case			)

   if writer._num_bytes > 0:
    UpperCamelCase_, UpperCamelCase_					:						str =					writer.finalize()
    writer.close()
    yield pa.RecordBatch.from_arrays(
        [[task_id], [num_examples], [num_bytes]]   ,      names=['task_id', 'num_examples', 'num_bytes']   ,      )

   if working_fpath != fpath:
    for file in os.listdir(os.path.dirname(snake_case			)			):
     UpperCamelCase_					:						Dict =					os.path.join(os.path.dirname(snake_case			)   ,      os.path.basename(snake_case			)			)
     shutil.move(snake_case   ,      snake_case			)

  UpperCamelCase_					:						int =					(
      self.df.mapInArrow(snake_case   ,      'task_id: long, num_examples: long, num_bytes: long'			)
      .groupBy('task_id'			)
      .agg(
          pyspark.sql.functions.sum('num_examples'			).alias('total_num_examples'			)   ,      pyspark.sql.functions.sum('num_bytes'			).alias('total_num_bytes'			)   ,      pyspark.sql.functions.count('num_bytes'			).alias('num_shards'			)   ,      pyspark.sql.functions.collect_list('num_examples'			).alias('shard_lengths'			)   ,      )
      .collect()
  )
  for row in stats:
   yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int   ,      snake_case	:							"datasets.SplitGenerator"   ,      snake_case	:							str = "arrow"   ,      snake_case	:							Optional[Union[str, int]] = None   ,      snake_case	:							Optional[int] = None   ,      **snake_case	:							Any   ,      )   ->						int:

  """simple docstring"""





  self._validate_cache_dir()

  UpperCamelCase_					:						Optional[int] =					convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE			)
  self._repartition_df_if_needed(snake_case			)
  UpperCamelCase_					:						List[str] =					not is_remote_filesystem(self._fs			)
  UpperCamelCase_					:						List[Any] =					os.path.join if is_local else posixpath.join

  UpperCamelCase_					:						Optional[int] =					'-TTTTT-SSSSS-of-NNNNN'
  UpperCamelCase_					:						Dict =					f"{self.name}-{split_generator.name}{SUFFIX}.{file_format}"
  UpperCamelCase_					:						int =					path_join(self._output_dir   ,      snake_case			)

  UpperCamelCase_					:						int =					0
  UpperCamelCase_					:						Optional[int] =					0
  UpperCamelCase_					:						Union[str, Any] =					0
  UpperCamelCase_					:						Optional[Any] =					[]
  UpperCamelCase_					:						Any =					[]

  for task_id, content in self._prepare_split_single(snake_case   ,      snake_case   ,      snake_case			):
   (
       (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), 
   )					:						Optional[Any] =					content
   if num_bytes > 0:
    total_num_examples += num_examples
    total_num_bytes += num_bytes
    total_shards += num_shards
    task_id_and_num_shards.append((task_id, num_shards)			)
    all_shard_lengths.extend(snake_case			)

  UpperCamelCase_					:						Optional[Any] =					total_num_examples
  UpperCamelCase_					:						Any =					total_num_bytes

  # should rename everything at the end
  logger.debug(f"Renaming {total_shards} shards."			)
  if total_shards > 1:
   UpperCamelCase_					:						List[Any] =					all_shard_lengths

   # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
   # pickling error due to pickling the SparkContext.
   UpperCamelCase_					:						int =					self._fs

   # use the -SSSSS-of-NNNNN pattern
   def _rename_shard(
       snake_case	:							int   ,      snake_case	:							int   ,      snake_case	:							int   ,      ):
    rename(
        snake_case   ,      fpath.replace('SSSSS'   ,      f"{shard_id:05d}"			).replace('TTTTT'   ,      f"{task_id:05d}"			)   ,      fpath.replace('TTTTT-SSSSS'   ,      f"{global_shard_id:05d}"			).replace('NNNNN'   ,      f"{total_shards:05d}"			)   ,      )

   UpperCamelCase_					:						Any =					[]
   UpperCamelCase_					:						Optional[int] =					0
   for i in range(len(snake_case			)			):
    UpperCamelCase_, UpperCamelCase_					:						Union[str, Any] =					task_id_and_num_shards[i]
    for shard_id in range(snake_case			):
     args.append([task_id, shard_id, global_shard_id]			)
     global_shard_id += 1
   self._spark.sparkContext.parallelize(snake_case   ,      len(snake_case			)			).map(lambda snake_case			: _rename_shard(*snake_case			)			).collect()
  else:
   # don't use any pattern
   UpperCamelCase_					:						Tuple =					0
   UpperCamelCase_					:						Optional[Any] =					task_id_and_num_shards[0][0]
   self._rename(
       fpath.replace('SSSSS'   ,      f"{shard_id:05d}"			).replace('TTTTT'   ,      f"{task_id:05d}"			)   ,      fpath.replace(snake_case   ,      ''			)   ,      )


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Tuple   ,      snake_case	:							"datasets.SplitGenerator"   ,      )   ->						SparkExamplesIterable:

  """simple docstring"""





  return SparkExamplesIterable(self.df			)



